Gemma-4-12B-PsiLM: bridges, Burgers FNO, inference script, model card

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+psilm-banner.png filter=lfs diff=lfs merge=lfs -text

MANIFEST.md

@@ -0,0 +1,77 @@
+# MANIFEST — Hugging Face repo `ryoji-info/Gemma-4-12B-PsiLM`
+
+Everything in this directory (`release/gemma-4-12b-psilm/`) is uploaded as-is to the
+model repo; nothing outside it is needed. The backbone (`mlx-community/gemma-4-12B-it-4bit`)
+is **not** uploaded — the script downloads it at first run.
+
+## Files to upload (and where each one came from)
+
+| path in the HF repo | source in the GitHub checkout | size | sha256 |
+|---|---|---:|---|
+| `README.md` | written for this release (the HF model card, YAML front matter) | 15 KB | — |
+| `psilm_infer.py` | written for this release | 18 KB | — |
+| `requirements.txt` | written for this release | 1 KB | — |
+| `MANIFEST.md` | this file (harmless to upload; drop it if you prefer) | — | — |
+| `psilm-banner.png` | `assets/psilm-banner.png` (1600 px wide) | 1.35 MB | `4cd65f32fe4aab5b49b66681a841b0128628e8d94c75494ae84f941b49c48089` |
+| `bridges/gemma-4-12b-4bit-mlx-1d-value-selective/bridges.safetensors` | `results/hf_export/bridges/gemma-4-12b-4bit-mlx-1d-value-selective/bridges.safetensors` (= `results/stage2_gemma12b/bridges.npz`, step 7000, learned-pointer tensors dropped) | 102,068,660 B | `f6ef8946c41b3cfa17df7c22bcab2c5cffbd856c2dd608970a364cc14d5e8f7d` |
+| `bridges/gemma-4-12b-4bit-mlx-1d-value-selective/config.json` | `results/hf_export/bridges/gemma-4-12b-4bit-mlx-1d-value-selective/config.json` | 1,287 B | `8a5add382beb65a0214f5e8b0b640e5d5df6f19e9a6caf1945400d2fe1e6bba5` |
+| `physics/fno_burgers_singlemode.safetensors` | `results/hf_export/physics/fno_burgers_singlemode.safetensors` (= `results/stage2/fno.pt`; loaders verified identical, max weight and field difference 0.0) | 552,076 B | `7bb0076c85cdcf2505a9079c05964e3eb77ac4a776eccf34216953f3c37bfcdd` |
+
+The copies were made with `cp` on 2026-09-06 and the hashes match the sources
+(`shasum -a 256`). Not uploaded: `__pycache__/`, `.DS_Store`.
+
+### To add later (the two runs in progress)
+
+When the Gemma multi-mode and 2D runs finish, export them the same way and add:
+
+| path in the HF repo | source |
+|---|---|
+| `bridges/gemma-4-12b-4bit-mlx-multimode/{bridges.safetensors,config.json}` | the stage-2b Gemma checkpoint (`results/stage2b_gemma12b*/bridges.npz` + a `config.json` in the format of the 1D one) |
+| `physics/fno_burgers_multimode.safetensors` | `results/hf_export/physics/fno_burgers_multimode.safetensors` (loads with `load_fno_safetensors`, verified) |
+| `bridges/gemma-4-12b-4bit-mlx-2d-dpot/{bridges.safetensors,config.json}` | the stage-2d Gemma checkpoint |
+| `physics/dpot_tiny_fisher2d_finetuned.safetensors` | `results/hf_export/physics/dpot_tiny_fisher2d_finetuned.safetensors` |
+
+and fill the two *in progress* rows of `README.md` ("Bridges in this repository") from their
+`final_eval.json`; `psilm_infer.py` will need a `--task` switch for those (it is 1D-only today).
+
+## Upload (maintainer runs this)
+
+From the GitHub checkout root, with `HF_TOKEN` set (or `huggingface-cli login` done):
+
+```bash
+cd /Users/rxiii/Documents/GitHub/PsiLM
+HF_HUB_DISABLE_XET=1 .venv/bin/python - <<'EOF'
+from huggingface_hub import HfApi
+api = HfApi()
+repo = "ryoji-info/Gemma-4-12B-PsiLM"
+api.create_repo(repo, repo_type="model", private=True, exist_ok=True)
+api.upload_folder(
+    folder_path="release/gemma-4-12b-psilm",
+    repo_id=repo,
+    repo_type="model",
+    commit_message="Gemma-4-12B-PsiLM: bridges, Burgers FNO, one-command CLI, model card",
+    ignore_patterns=["__pycache__/*", "*.pyc", ".DS_Store"],
+)
+print("UPLOAD DONE")
+EOF
+```
+
+`HF_HUB_DISABLE_XET=1` keeps the upload on the classic LFS path (the Xet backend has
+stalled on this machine's earlier uploads). Make the repo public afterwards with
+`api.update_repo_settings(repo, private=False)` or from the repo's settings page, together
+with `ryoji-info/PsiLM-bridges` and `ryoji-info/PsiLM-physics`, which the card links to.
+
+## After upload: the one-command check
+
+```bash
+huggingface-cli download ryoji-info/Gemma-4-12B-PsiLM --local-dir /tmp/g4psilm && cd /tmp/g4psilm
+python3 -m venv .venv && .venv/bin/pip install -r requirements.txt
+.venv/bin/python psilm_infer.py
+```
+
+Note for the `pip install` line: the `psilm @ git+...` requirement needs the GitHub
+`pyproject.toml` to use package discovery (`[tool.setuptools.packages.find] include = ["psilm*"]`,
+changed on 2026-09-06 in the working tree from `packages = ["psilm"]`, which shipped only the
+top-level package without `psilm.mlx`, `psilm.stage2`, `psilm.physics`). That change must be
+pushed before the requirement installs a usable package; until then the script says so, and the
+`PSILM_REPO=/path/to/clone` fallback works regardless.

README.md

@@ -0,0 +1,182 @@
+---
+license: apache-2.0
+base_model: mlx-community/gemma-4-12B-it-4bit
+library_name: mlx
+tags:
+  - mlx
+  - psilm
+  - physics
+  - gemma
+---
+
+![PsiLM](psilm-banner.png)
+
+# Gemma-4-12B-PsiLM
+
+**Gemma 4 12B, coupled to a physics model through trained latent bridges, runnable on a Mac with one command.**
+
+*Research generated by Claude Fable 5 (Anthropic) under the direction of Ryoji Furui; see the [AI generation disclosure](#ai-generation-disclosure).*
+
+## What this is
+
+PsiLM (ΨLM) runs a **frozen language model** and a **frozen physics model** *together* while one answer is produced. Nothing is fine-tuned and no text crosses the interface: small trainable **bridges** read the physical problem out of the language model's hidden states, the physics model computes, and its result flows back into the language model's residual stream as a few soft tokens through a gated cross-attention. This repository packages that system for one backbone:
+
+| part | what | trained? | where it comes from |
+|---|---|---|---|
+| language model | Gemma 4 12B-it, 4-bit MLX quantization (48 layers, hidden 3840) | frozen | downloaded from [`mlx-community/gemma-4-12B-it-4bit`](https://huggingface.co/mlx-community/gemma-4-12B-it-4bit) at first run (not redistributed here) |
+| physics model | a 1D Burgers Fourier Neural Operator, 70K parameters (`physics/fno_burgers_singlemode.safetensors`) | frozen | this repo (also in [`ryoji-info/PsiLM-physics`](https://huggingface.co/ryoji-info/PsiLM-physics)) |
+| bridges | forward readout + value-token channel + gated injection, 25.5M parameters (`bridges/gemma-4-12b-4bit-mlx-1d-value-selective/`) | **trained** | this repo (also in [`ryoji-info/PsiLM-bridges`](https://huggingface.co/ryoji-info/PsiLM-bridges)) |
+
+**What it can answer today.** One family of questions, the one the bridges were trained on:
+
+> A velocity field on the periodic domain [0,1) starts as u(x,0) = *a* · sin(2πx + *φ*). It evolves by Burgers' equation with viscosity 0.02 until t = 0.5. What is the value of u at x = *x₀*? Answer with a number rounded to 2 decimal places.
+
+with *a* in [0.5, 1.5], *φ* in [0, 6.28], *x₀* in [0, 0.99], two decimals each. Gemma alone cannot answer these (0% on the held-out set; it derives for 768 tokens and never commits to a number). Coupled through the bridges it answers **96.7%** of them within ±0.05, against **98.3%** when the true value is written into the prompt as text (the oracle ceiling). On everything else the gate stays shut and the model is Gemma, byte for byte in the arms measured (GSM8K 84/84, see below).
+
+**What it is not.** Not a general physics assistant, not a fine-tuned Gemma, and not a model that will know when a *different* PDE applies: the physics model solves exactly one equation family and the bridges read exactly three quantities (*a*, *φ*, *x₀*) from the text. It is a research artifact: a working, measured instance of latent coupling between a language model and a physics model at 12B scale, on consumer hardware.
+
+## Run it
+
+Apple Silicon Mac (the 4-bit backbone is 6.3 GB on disk; training of the bridges peaked at 13 GB on a 24 GB M2, inference needs less), Python 3.11+.
+
+```bash
+# 1. get this repo (≈100 MB of bridges + 0.5 MB of physics model)
+huggingface-cli download ryoji-info/Gemma-4-12B-PsiLM --local-dir Gemma-4-12B-PsiLM
+cd Gemma-4-12B-PsiLM
+
+# 2. dependencies (mlx, mlx-lm, transformers, torch, huggingface_hub + the psilm package from GitHub)
+pip install -r requirements.txt
+
+# 3. one command
+python psilm_infer.py
+```
+
+The first run downloads the Gemma 4 backbone (6.3 GB) into the Hugging Face cache. The script answers the default question (*a* = 1.28, *φ* = 0.5, *x₀* = 0.76) three ways and prints timing:
+
+```
+[1] PsiLM (coupled)     : the coupled system's reply ('u at x = 0.76 equals <number>.')   -- seconds
+[2] backbone alone      : Gemma 4 alone under the "Answer: <number>" protocol, answer-forced -- about 75 s
+[3] physics model (FNO) : u(0.76) = -0.2522   (the FNO on the true initial condition: the reference [1] should match to ±0.05)
+    spectral solver     : u(0.76) = -0.2517   (ground truth)
+```
+
+Other questions and options:
+
+```bash
+python psilm_infer.py --a 0.9 --phi 2.1 --x0 0.33     # any (a, phi, x0) in the ranges above
+python psilm_infer.py --no-baseline                   # skip the slow backbone-alone arm
+python psilm_infer.py --question-only                 # print the exact prompt, load nothing
+python psilm_infer.py --help                          # --bridges DIR, --physics FILE, --backbone ID, ...
+```
+
+Without `pip install`-ing the package, a clone of the GitHub repository works too: `PSILM_REPO=/path/to/PsiLM python psilm_infer.py`.
+
+Loading the pieces yourself, in Python:
+
+```python
+import json, mlx.core as mx
+from psilm.mlx.gemma_loader import load_backbone_any     # Gemma 4 text tower in the staged-forward layout
+from psilm.mlx.bridges import PsiBridgesMLX
+from psilm.mlx.fno import load_fno_safetensors
+from psilm.mlx.model import PsiLMMLX
+
+model, stock, tok = load_backbone_any("mlx-community/gemma-4-12B-it-4bit")
+d = "bridges/gemma-4-12b-4bit-mlx-1d-value-selective"
+cfg = json.load(open(f"{d}/config.json"))
+bridges = PsiBridgesMLX(**cfg["construct"])                # d_model 3840, channel "value", inj_cap 0.2, readout_norm "dim"
+bridges.load_weights(f"{d}/bridges.safetensors", strict=False)   # the retired learned-pointer tensors are omitted
+fno = load_fno_safetensors("physics/fno_burgers_singlemode.safetensors")
+psi = PsiLMMLX(model, tok, fno, bridges, l_fwd=cfg["coupling"]["l_fwd"], l_rev=cfg["coupling"]["l_rev"])   # read @20, inject @30 of 48
+# psi.generate(QABuilder(hf_tokenizer), {"a": 1.28, "phi": 0.5, "x0": 0.76}) -- see psilm_infer.py
+```
+
+## Results for this backbone
+
+Held-out evaluation, 60 questions, accuracy within ±0.05 (`results/stage2_gemma12b/final_eval.json` in the [GitHub repository](https://github.com/ryoji-info/PsiLM)):
+
+| arm | accuracy | MAE | note |
+|---|---:|---:|---|
+| Gemma 4 12B alone | 0.0% | 2.93 | never reaches an `Answer:` line within 768 tokens |
+| **PsiLM (this repo)** | **96.7%** | **0.017** | bridges read the prompt, FNO computes, value returns in latent space |
+| oracle (true value written into the prompt) | 98.3% | 0.007 | the tool-loop ceiling |
+| always answer 0.00 | 1.7% | 0.308 | calibration |
+
+Guard-rail: does the coupled model still do everything else? 100 questions per dataset, three arms — backbone alone / PsiLM / PsiLM with the injection zeroed — with the gate recorded per question (`results/bench/gemma12b_guardrail_summary.json`, `results/bench/gemma12b_nonudge_guardrail_summary.json`):
+
+| dataset (n=100) | backbone | **PsiLM** | zeroed | gate σ (PsiLM) | open on |
+|---|---:|---:|---:|---:|---:|
+| physics QA (this task) | 0% | **97%** | 10% | 0.14 | 100% |
+| GSM8K | 84% | **84%** | 84% | 0.004 | 0% |
+| GSM8K, no `Answer:` line in the prompt | 83% | **83%** | 83% | 0.002 | 0% |
+| MMLU, 5 subjects, 256 tokens | 53% | **55%** | 53% | 0.008 | 0% |
+
+On MMLU the two arms agree at 79.1% / 79.1% on the 67 items both answer within the budget; the raw 53 vs 55 is parse noise. The gate selectivity comes from a **no-harm training arm**: 1,046 non-physics prompts (GSM8K train, MMLU validation, with and without the `Answer:` nudge) paired with Gemma's own greedy continuations, on which only the gate is updated with a mean-gate penalty (`config.json` → `training`). The 10% of the zeroed arm on physics is what the reply template alone recovers; the coupled 97% is against that floor.
+
+Training: 7,000 steps at batch 4 on one Apple M2 (24 GB) — 2,000 readout-only warm-up steps, 3,500 coupled steps (lr 3e-4; the coupled phase ran on to step 6,000 but the no-harm phase resumed from the step-5,500 checkpoint), then 1,500 no-harm steps at lr 1e-4 (`results/gemma12b/noharm_recipe.sh`); 12 s per step at a 13 GB peak. One backbone-specific adjustment, measured rather than tuned: Gemma's massive-activation dimensions are nearly constant across prompts, so the readout standardizes each hidden dimension with statistics from a 32-prompt calibration pass (`readout_norm: "dim"`; the two frozen vectors `fwd.dim_mu`, `fwd.dim_sigma` are in the checkpoint).
+
+### Bridges in this repository
+
+| directory | task | physics model | trained params | held-out | status |
+|---|---|---|---|---:|---|
+| `bridges/gemma-4-12b-4bit-mlx-1d-value-selective/` | 1D Burgers, single-mode initial conditions, value at *x₀* | `physics/fno_burgers_singlemode.safetensors` | 25.5M | **96.7%** @±0.05, MAE 0.017 (n=60); GSM8K 84/84 | **released** |
+| `bridges/gemma-4-12b-4bit-mlx-multimode/` | 1D Burgers, multi-mode initial conditions (modes 1–4, mixed amplitudes) + generalization families | `physics/fno_burgers_multimode.safetensors` | — | *in progress* | run in progress; row filled when the evaluation lands |
+| `bridges/gemma-4-12b-4bit-mlx-2d-dpot/` | 2D Fisher–KPP, replicated-IC history → u(0.4), value at (x₀, y₀) | `physics/dpot_tiny_fisher2d_finetuned.safetensors` (DPOT-Tiny, 7.5M, fine-tuned) | — | *in progress* | run in progress; row filled when the evaluation lands |
+
+The two in-progress rows are the multi-mode (stage 2b) and 2D (stage 2d) tasks that the 0.5B backbone completed earlier (iid 97.9% and 95.0% respectively, see [`ryoji-info/PsiLM-bridges`](https://huggingface.co/ryoji-info/PsiLM-bridges)); their Gemma runs are training now and will be added with their own `config.json`, the matching physics file, and a `--task` switch in `psilm_infer.py`.
+
+## How it works, in one paragraph
+
+The prompt runs through Gemma's first 20 layers. The **forward bridge** reads the queried position *x₀* by pooling the hidden states over its tokens (a deterministic span pointer computed by the QA builder, plus a 100-bin classifier) and the initial-condition parameters with a learned pool, after the calibrated per-dimension standardization) and emits (*a*, sin *φ*, cos *φ*) and *x₀*; from these it builds the initial condition on a 128-point grid. The frozen **FNO** evolves it to t = 0.5. A learned periodic lookup kernel reads the field at *x₀*, and the **value-token channel** turns that single number into eight soft tokens through Fourier features. At layer 30 a **gated cross-attention** injects them into the residual stream, capped at 20% of the stream's RMS; the gate is a small MLP on the residual stream, trained to open on physics prompts and close elsewhere. Layers 30–48 and the answer are Gemma's own. Details, ablations and the failure analysis that produced this design are in the paper (`paper/psilm.pdf` in the repository, Section 9 for scaling, the guard-rail and Gemma).
+
+## Limitations
+
+- **One task family.** The bridges read exactly the three quantities of the trained question and the FNO solves exactly one equation family; a different PDE, boundary condition, viscosity or final time is out of scope, and the gate closing on non-physics text does not mean it can recognize *other* physics. Free-text initial conditions ("a Gaussian bump near the left edge") are not supported.
+- **The pointer is task-supplied.** Which tokens hold *x₀* is computed by the QA builder from the prompt (`QABuilder.x0_span`), not learned from the words; the learned attention pointer never left uniform at 8B and was retired. `psilm_infer.py` builds the prompt itself for that reason; a paraphrased question is not the trained input.
+- **4-bit backbone, quantized kernels.** The bridges were trained *through* the 4-bit backbone and absorb its quantization noise, but MLX's quantized matmuls are not guaranteed bit-identical across `mlx` versions or Apple chips, so a given question can land a hundredth away from the recorded run; the numbers above are from mlx 0.32.2 / mlx-lm 0.31.3 on an M2.
+- **Inputs with two decimals, inside the training ranges.** The readout was trained on numbers formatted like the training set; `psilm_infer.py` rounds inputs to two decimals and warns outside [0.5, 1.5] × [0, 6.28].
+- **Bridges do not transfer between backbones.** This checkpoint is for `mlx-community/gemma-4-12B-it-4bit` exactly (hidden 3840, 48 layers); the script refuses a backbone of another width.
+- **Sequence length.** The Gemma loader drives the sliding-window layers with a plain causal mask, exact up to the 1024-token window; the questions here are ~120 tokens.
+- **Evaluation scope.** The guard-rail covers GSM8K, a five-subject MMLU slice and this physics set at n=100 each; nothing else has been measured.
+
+## Beyond physics
+
+The bridges here couple a frozen language model to a frozen *physics* model, but the recipe (read a fixed set of quantities from text; let a frozen quantitative model compute; return one value through a selective gate) is not specific to PDEs. A calibrated market or event-probability model in the physics model's seat would be the same architecture, and the appeal is the same: a language model's forecast grounded in a model that can be validated separately, with a gate that stays shut when the model does not apply. Nothing in this repository has been trained or tested on financial data; the physics results relied on exact oracles, deterministic targets and no distribution shift, none of which markets provide. This is a research direction, not a capability, and not a basis for investment decisions.
+
+## Files
+
+```
+psilm_infer.py          the one-command CLI (PsiLM / backbone alone / physics model)
+requirements.txt        pip dependencies, including the psilm package from GitHub
+bridges/gemma-4-12b-4bit-mlx-1d-value-selective/
+    bridges.safetensors     the trained bridges (25.5M params, fp32, 97 MB)
+    config.json             construction, coupling depths, training record, per-chunk held-out scores
+physics/fno_burgers_singlemode.safetensors   the frozen FNO (torch key names; complex spectral weights as .real/.imag)
+psilm-banner.png        banner
+README.md               this card
+```
+
+## Related
+
+- Code, training scripts, evaluation records and the paper: [github.com/ryoji-info/PsiLM](https://github.com/ryoji-info/PsiLM)
+- Bridges for every backbone (Qwen2.5-0.5B, Qwen3-1.7B, Qwen3-8B, Gemma 4 12B): [ryoji-info/PsiLM-bridges](https://huggingface.co/ryoji-info/PsiLM-bridges)
+- The frozen physics models: [ryoji-info/PsiLM-physics](https://huggingface.co/ryoji-info/PsiLM-physics)
+
+## Citation
+
+```bibtex
+@misc{furui2026psilm,
+  title  = {PsiLM: Coupling Frozen Language and Physics Models through Trainable Latent Bridges},
+  author = {Furui, Ryoji},
+  year   = {2026},
+  url    = {https://github.com/ryoji-info/PsiLM},
+  note   = {Research generated by Claude Fable 5 (Anthropic) under the author's direction}
+}
+```
+
+## License
+
+The bridges, the FNO and the code in this repository are released under **Apache-2.0**. The Gemma 4 backbone is not part of this repository: `psilm_infer.py` downloads it from `mlx-community/gemma-4-12B-it-4bit`, and its use is governed by Google's [Gemma Terms of Use](https://ai.google.dev/gemma/terms).
+
+## AI generation disclosure
+
+This model, its training recipe, the evaluations and this card were generated by **Claude Fable 5** (Anthropic), operating as an autonomous research agent under the direction and review of Ryoji Furui, who set the research question and the hardware constraint, approved each stage, and bears responsibility for the published claims. All numbers on this card are taken from committed evaluation records in the repository, cited by file name above.

bridges/gemma-4-12b-4bit-mlx-1d-value-selective/bridges.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6ef8946c41b3cfa17df7c22bcab2c5cffbd856c2dd608970a364cc14d5e8f7d
+size 102068660

bridges/gemma-4-12b-4bit-mlx-1d-value-selective/config.json

@@ -0,0 +1,47 @@
+{
+ "backbone": "mlx-community/gemma-4-12B-it-4bit",
+ "hf_tokenizer": "mlx-community/gemma-4-12B-it-4bit",
+ "loader": "psilm.mlx.gemma_loader.load_backbone_any (gemma4_unified -> text tower)",
+ "physics": "results/stage2/fno.pt (1D Burgers FNO, 70K params)",
+ "bridges_class": "psilm.mlx.bridges.PsiBridgesMLX",
+ "construct": {
+  "d_model": 3840,
+  "channel": "value",
+  "inj_cap": 0.2,
+  "gate_bias": 0.0,
+  "readout_norm": "dim"
+ },
+ "coupling": {
+  "l_fwd": 20,
+  "l_rev": 30,
+  "n_layers": 48
+ },
+ "pointer": "deterministic span pooling with calibrated per-dimension standardization (fwd.dim_mu/dim_sigma included; fwd.x0_query/x0_key omitted: unused)",
+ "training": {
+  "steps_total": 7000,
+  "phase_A_readout_only": 2000,
+  "coupled_steps": 4000,
+  "no_harm_steps": 1500,
+  "batch": 4,
+  "lr": "3e-4 (warm-up, coupled), 1e-4 (no-harm phase, from the step-5500 checkpoint)",
+  "no_harm_arm": "1046 non-physics prompts (GSM8K train, MMLU validation; with/without the Answer nudge) paired with the backbone's own greedy continuation; gate-only updates + mean-gate penalty"
+ },
+ "held_out_n48_per_chunk": {
+  "coupled": [
+   0.625,
+   0.667,
+   0.625,
+   0.875,
+   0.875,
+   0.625,
+   0.792,
+   0.542
+  ],
+  "no_harm_phase": [
+   1.0,
+   1.0,
+   0.979
+  ]
+ },
+ "license": "apache-2.0"
+}

physics/fno_burgers_singlemode.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7bb0076c85cdcf2505a9079c05964e3eb77ac4a776eccf34216953f3c37bfcdd
+size 552076

psilm_infer.py

@@ -0,0 +1,364 @@
+#!/usr/bin/env python3
+"""PsiLM on Gemma 4 12B: one command, three numbers.
+
+    python psilm_infer.py                          # a=1.28, phi=0.5, x0=0.76
+    python psilm_infer.py --a 0.9 --phi 2.1 --x0 0.33
+    python psilm_infer.py --question-only          # print the prompt, load nothing
+    python psilm_infer.py --no-baseline            # skip the (slow) backbone-alone arm
+
+Loads a frozen 4-bit language model (default ``mlx-community/gemma-4-12B-it-4bit``),
+the trained PsiLM bridges for it (``bridges/<name>/{bridges.safetensors,config.json}``)
+and the frozen 1D Burgers FNO (``physics/fno_burgers_singlemode.safetensors``), then
+answers one Burgers field-value question three ways:
+
+  PsiLM     the coupled system: the forward bridge reads the initial condition and
+            the queried position x0 out of the prompt's hidden states, the FNO evolves
+            the field, and the looked-up value u(x0) returns to the language model as
+            soft tokens through a gated cross-attention. No text crosses the interface.
+  backbone  the same language model alone, "Answer: <number>" protocol with answer
+            forcing (the baseline arm of eval/mlx_stage2_eval.py in the repository).
+  physics   the FNO's own value at x0 on the TRUE initial condition -- the reference
+            the coupled answer should match to +-0.05 -- plus the spectral solver's
+            ground truth for the same question.
+
+Dependencies: ``pip install -r requirements.txt`` (mlx, mlx-lm, transformers, torch,
+huggingface_hub, and the ``psilm`` package from https://github.com/ryoji-info/PsiLM).
+If ``psilm`` is not installed, set ``PSILM_REPO=/path/to/PsiLM`` (a clone) instead.
+"""
+
+import argparse
+import importlib
+import importlib.util
+import json
+import math
+import os
+import re
+import sys
+import time
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+DEFAULT_BACKBONE = "mlx-community/gemma-4-12B-it-4bit"
+DEFAULT_BRIDGES = HERE / "bridges" / "gemma-4-12b-4bit-mlx-1d-value-selective"
+DEFAULT_PHYSICS = HERE / "physics" / "fno_burgers_singlemode.safetensors"
+TOL = 0.05
+
+
+# --------------------------------------------------------------------------- imports
+def _psilm_available():
+    try:
+        return importlib.util.find_spec("psilm.mlx") is not None
+    except ModuleNotFoundError:
+        return False
+
+
+def _ensure_psilm():
+    """Import path for the ``psilm`` package: the installed package first, else the
+    optional ``PSILM_REPO`` environment variable (a clone of the GitHub repository)."""
+    if _psilm_available():
+        return None
+    repo = os.environ.get("PSILM_REPO")
+    if repo and (Path(repo) / "psilm" / "mlx").is_dir():
+        sys.path.insert(0, str(Path(repo).resolve()))
+        for name in [m for m in sys.modules if m == "psilm" or m.startswith("psilm.")]:
+            del sys.modules[name]                 # a partial install may already be imported
+        importlib.invalidate_caches()
+        if _psilm_available():
+            return str(Path(repo).resolve())
+    sys.exit(
+        "psilm_infer.py: the 'psilm' package (with its psilm.mlx subpackage) is not importable.\n"
+        "  Install it:   pip install -r requirements.txt\n"
+        "                (or: pip install 'git+https://github.com/ryoji-info/PsiLM')\n"
+        "  Or point at a clone:  PSILM_REPO=/path/to/PsiLM python psilm_infer.py ..."
+    )
+
+
+PSILM_REPO = _ensure_psilm()
+
+import mlx.core as mx  # noqa: E402
+from mlx.utils import tree_flatten  # noqa: E402
+import mlx_lm  # noqa: E402
+import numpy as np  # noqa: E402
+from transformers import AutoTokenizer  # noqa: E402
+
+from psilm.mlx.bridges import PsiBridgesMLX, build_ic_mlx  # noqa: E402
+from psilm.mlx.fno import convert_from_torch, load_fno_safetensors  # noqa: E402
+from psilm.mlx.gemma_loader import load_backbone_any  # noqa: E402
+from psilm.mlx.model import PsiLMMLX  # noqa: E402
+from psilm.physics.burgers import initial_condition, solve  # noqa: E402
+from psilm.stage2.qa import QUESTION, SYSTEM, QABuilder, fourier_interp  # noqa: E402
+
+
+# --------------------------------------------------------------- backbone-alone arm
+# The baseline protocol is eval/mlx_stage2_eval.py's: the question plus the "Answer:"
+# nudge, greedy generation, and answer forcing when the reply used its budget without
+# an Answer line. When a repository clone is reachable (PSILM_REPO) the functions are
+# imported from that file so the release cannot drift from the evaluation; otherwise
+# the verbatim copies below are used.
+NUDGE = "\nEnd your reply with a line of the form \"Answer: <number>\"."
+FORCE_SUFFIX = "\n\nAnswer:"
+
+
+def parse_value(text):
+    m = re.findall(r"Answer:\s*\$?\\?\(?\s*(-?\d+\.?\d*)", text)
+    if m:
+        return float(m[-1])
+    m = re.findall(r"-?\d+\.\d+", text)
+    return float(m[-1]) if m else None
+
+
+def chat_generate(model, hf_tok, user, max_new=768, gen_tok=None, force_answer=True):
+    """Returns (text, forced). gen_tok is the mlx-lm tokenizer wrapper (knows all of
+    a backbone's stop ids, e.g. Gemma's <eos>/<turn|>); hf_tok builds the prompt."""
+    messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": user}]
+    ids = hf_tok.apply_chat_template(messages, tokenize=True, add_generation_prompt=True,
+                                     enable_thinking=False)
+    if not isinstance(ids, list):
+        ids = ids["input_ids"]
+    if ids and isinstance(ids[0], list):
+        ids = ids[0]
+    text = mlx_lm.generate(model, gen_tok or hf_tok, prompt=list(ids), max_tokens=max_new, verbose=False)
+    if force_answer and "Answer:" not in text:
+        cont = list(ids) + hf_tok.encode(text + FORCE_SUFFIX, add_special_tokens=False)
+        tail = mlx_lm.generate(model, gen_tok or hf_tok, prompt=cont, max_tokens=16, verbose=False)
+        return text + FORCE_SUFFIX + tail, True
+    return text, False
+
+
+def _baseline_protocol():
+    """(chat_generate, parse_value, NUDGE, source): the repository's eval functions
+    when a clone is reachable, else the copies in this file."""
+    roots = [Path(PSILM_REPO)] if PSILM_REPO else []
+    for root in roots:
+        f = root / "eval" / "mlx_stage2_eval.py"
+        if f.is_file():
+            spec = importlib.util.spec_from_file_location("psilm_release_eval_stage2", f)
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+            return mod.chat_generate, mod.parse_value, mod.NUDGE, str(f)
+    return chat_generate, parse_value, NUDGE, "vendored copy (eval/mlx_stage2_eval.py)"
+
+
+def parse_psilm_answer(text):
+    """The trained reply is 'u at x = {x0} equals {u}.'; the number after 'equals' is
+    the answer (the evaluation scores only that). Returns (value, strict): strict is
+    False when the reply left the template and the last decimal number is used."""
+    m = re.search(r"equals\s*(-?\d+\.?\d*)", text)
+    if m:
+        return float(m.group(1)), True
+    m = re.findall(r"-?\d+\.\d+", text)
+    return (float(m[-1]), False) if m else (None, False)
+
+
+# ------------------------------------------------------------------------ bridges
+def load_bridges(bridges_dir, d_model, n_layers, args):
+    """PsiBridgesMLX from <dir>/config.json['construct'] (defaults when the file is
+    absent, inferred from the tensors present) + <dir>/bridges.safetensors with
+    load_weights(strict=False): the Gemma export omits the retired learned-pointer
+    tensors (fwd.x0_query, fwd.x0_key.*), which the deterministic span pointer never
+    uses. Shapes of every provided tensor are checked explicitly, since non-strict
+    loading would not. Returns (bridges, config, coupling, report)."""
+    bridges_dir = Path(bridges_dir)
+    cfg_path = bridges_dir / "config.json"
+    cfg = json.loads(cfg_path.read_text()) if cfg_path.is_file() else {}
+    weights = mx.load(str(bridges_dir / "bridges.safetensors"))
+    keys = set(weights)
+
+    construct = {
+        "d_model": int(weights["fwd.query"].shape[0]) if "fwd.query" in keys else d_model,
+        "channel": "value" if any(k.startswith("val.") for k in keys) else "field",
+        "gate_bias": -2.0,
+        "inj_cap": None,
+        "readout_norm": "dim" if "fwd.dim_mu" in keys else "rms",
+    }
+    construct.update(cfg.get("construct", {}))
+    for name in ("channel", "gate_bias", "inj_cap", "readout_norm"):   # CLI overrides
+        v = getattr(args, name)
+        if v is not None:
+            construct[name] = v
+    if int(construct["d_model"]) != int(d_model):
+        sys.exit(f"bridges were trained for hidden size {construct['d_model']} but the "
+                 f"backbone has {d_model}: bridges do not transfer between backbones")
+    bridges = PsiBridgesMLX(**construct)
+
+    # explicit shape check + missing/unexpected report
+    params = dict(tree_flatten(bridges.parameters()))
+    missing = sorted(k for k in params if k not in keys)
+    unexpected = sorted(k for k in keys if k not in params)
+    bad = [(k, tuple(weights[k].shape), tuple(params[k].shape))
+           for k in keys if k in params and tuple(weights[k].shape) != tuple(params[k].shape)]
+    if bad:
+        lines = "\n".join(f"  {k}: file {a} vs module {b}" for k, a, b in bad)
+        sys.exit(f"bridge tensor shapes do not match the constructed bridges:\n{lines}")
+    if unexpected:
+        sys.exit(f"bridges.safetensors has tensors the bridges do not define: {unexpected}")
+    bridges.load_weights(list(weights.items()), strict=False)
+    mx.eval(bridges.parameters())
+
+    coupling = dict(cfg.get("coupling", {}))
+    if args.l_fwd is not None:
+        coupling["l_fwd"] = args.l_fwd
+    if args.l_rev is not None:
+        coupling["l_rev"] = args.l_rev
+    if "n_layers" in coupling and int(coupling["n_layers"]) != int(n_layers):
+        sys.exit(f"config.json expects a {coupling['n_layers']}-layer backbone; "
+                 f"this one has {n_layers}")
+    n_trained = sum(int(np.prod(v.shape)) for k, v in params.items() if k in keys)
+    report = {"construct": construct, "missing": missing, "n_tensors": len(keys),
+              "n_params": n_trained}
+    return bridges, cfg, coupling, report
+
+
+def load_physics(path):
+    path = str(path)
+    if path.endswith(".pt"):
+        return convert_from_torch(path)
+    return load_fno_safetensors(path)
+
+
+# --------------------------------------------------------------------------- main
+def build_parser():
+    ap = argparse.ArgumentParser(
+        description="PsiLM (frozen LLM + frozen Burgers FNO through latent bridges): "
+                    "answer one field-value question three ways.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="Training ranges: a in [0.5, 1.5], phi in [0, 6.28], x0 in [0, 0.99], "
+               "all with two decimals. Outside them the bridges are extrapolating.")
+    ap.add_argument("--a", type=float, default=1.28, help="amplitude of u(x,0) = a sin(2 pi x + phi)")
+    ap.add_argument("--phi", type=float, default=0.5, help="phase (radians)")
+    ap.add_argument("--x0", type=float, default=0.76, help="queried position in [0, 1)")
+    ap.add_argument("--backbone", default=DEFAULT_BACKBONE,
+                    help="mlx-lm checkpoint (Hugging Face id or local path)")
+    ap.add_argument("--bridges", default=str(DEFAULT_BRIDGES),
+                    help="directory with bridges.safetensors (+ config.json)")
+    ap.add_argument("--physics", default=str(DEFAULT_PHYSICS),
+                    help="FNO weights: safetensors export, or a PyTorch fno.pt")
+    ap.add_argument("--hf-tokenizer", default=None,
+                    help="HF tokenizer id for the chat template (default: config.json's "
+                         "hf_tokenizer, else the backbone id)")
+    ap.add_argument("--no-baseline", action="store_true", help="skip the backbone-alone arm")
+    ap.add_argument("--question-only", action="store_true",
+                    help="print the question the models see and exit (loads nothing)")
+    ap.add_argument("--max-new", type=int, default=24, help="PsiLM reply budget (tokens)")
+    ap.add_argument("--baseline-max-new", type=int, default=768,
+                    help="backbone-alone reply budget before answer forcing (the "
+                         "evaluation used 768; Gemma 4 uses all of it, ~75 s)")
+    g = ap.add_argument_group("bridge construction overrides (default: config.json)")
+    g.add_argument("--channel", choices=["value", "field"], default=None)
+    g.add_argument("--gate-bias", type=float, default=None)
+    g.add_argument("--inj-cap", type=float, default=None)
+    g.add_argument("--readout-norm", choices=["rms", "dim"], default=None)
+    g.add_argument("--l-fwd", type=int, default=None, help="readout layer")
+    g.add_argument("--l-rev", type=int, default=None, help="injection layer")
+    ap.add_argument("-v", "--verbose", action="store_true")
+    return ap
+
+
+def _clean(text):
+    return re.sub(r"<\|?[a-z_|]+\|?>", "", text).strip()
+
+
+def main():
+    args = build_parser().parse_args()
+    item = {"a": round(args.a, 2), "phi": round(args.phi, 2), "x0": round(args.x0, 2)}
+    if (item["a"], item["phi"], item["x0"]) != (args.a, args.phi, args.x0):
+        print(f"note: inputs rounded to two decimals (the trained readout reads 2-decimal "
+              f"numbers): a={item['a']} phi={item['phi']} x0={item['x0']}")
+    if not (0.0 <= item["x0"] < 1.0):
+        sys.exit("x0 must lie in [0, 1): the domain is periodic")
+    if not (0.5 <= item["a"] <= 1.5 and 0.0 <= item["phi"] <= 6.28):
+        print("warning: a or phi lies outside the training ranges (a in [0.5, 1.5], "
+              "phi in [0, 6.28]); the bridges are extrapolating")
+    question = QUESTION.format(a=item["a"], phi=item["phi"], x0=item["x0"])
+
+    if args.question_only:
+        print(f"[system] {SYSTEM}")
+        print(f"[user]   {question}")
+        print(f"\n(the backbone-alone arm appends: {NUDGE.strip()!r})")
+        return
+
+    # ---- load
+    t0 = time.time()
+    print(f"backbone : {args.backbone}", flush=True)
+    model, stock, tok = load_backbone_any(args.backbone)
+    n_layers = len(model.model.layers)
+    d_model = int(model.args.hidden_size)
+    bridges, cfg, coupling, rep = load_bridges(args.bridges, d_model, n_layers, args)
+    hf_id = args.hf_tokenizer or cfg.get("hf_tokenizer") or args.backbone
+    hf_tok = AutoTokenizer.from_pretrained(hf_id)
+    fno = load_physics(args.physics)
+    psi = PsiLMMLX(model, tok, fno, bridges, l_fwd=coupling.get("l_fwd"), l_rev=coupling.get("l_rev"))
+    builder = QABuilder(hf_tok)
+    t_load = time.time() - t0
+    c = rep["construct"]
+    print(f"bridges  : {args.bridges}\n"
+          f"           channel={c['channel']} readout_norm={c['readout_norm']} "
+          f"gate_bias={c['gate_bias']} inj_cap={c['inj_cap']} | "
+          f"{rep['n_params']/1e6:.1f}M params in {rep['n_tensors']} tensors"
+          + (f" | not in file (unused): {rep['missing']}" if rep["missing"] else ""))
+    print(f"coupling : read @ layer {psi.l_fwd}, inject @ layer {psi.l_rev} of {psi.n_layers} "
+          f"(hidden {d_model}"
+          + (")" if cfg.get("coupling") else "; no config.json coupling entry: PsiLMMLX defaults)"))
+    # parameter count with each complex spectral weight counted once (wr/wi are one number)
+    n_fno = sum(int(np.prod(v.shape)) for k, v in tree_flatten(fno.parameters()) if not k.endswith(".wi"))
+    print(f"physics  : {args.physics} (FNO1d, {n_fno/1e3:.0f}K params)")
+    print(f"loaded in {t_load:.1f} s\n")
+    print(f"question : {question}\n", flush=True)
+
+    # ---- 1. PsiLM: the coupled system
+    t1 = time.time()
+    psi_text = psi.generate(builder, item, max_new=args.max_new)
+    t_psi = time.time() - t1
+    psi_val, strict = parse_psilm_answer(psi_text)
+    print(f"[1] PsiLM (coupled)     : {_clean(psi_text)!r}")
+    print(f"    value               : {psi_val if psi_val is not None else 'no number parsed'}"
+          f"   ({t_psi:.1f} s{'' if strict or psi_val is None else '; reply left the trained template, last number taken'})",
+          flush=True)
+
+    # ---- 2. backbone alone
+    base_val = None
+    if not args.no_baseline:
+        gen, parse, nudge, src = _baseline_protocol()
+        t2 = time.time()
+        base_text, forced = gen(stock, hf_tok, question + nudge, args.baseline_max_new,
+                                gen_tok=tok, force_answer=True)
+        t_base = time.time() - t2
+        base_val = parse(base_text)
+        tail = _clean(base_text)
+        tail = tail if args.verbose or len(tail) <= 240 else "..." + tail[-240:]
+        print(f"[2] backbone alone      : {tail!r}")
+        print(f"    value               : {base_val if base_val is not None else 'no number parsed'}"
+              f"   ({t_base:.1f} s, {len(tok.encode(base_text))} tokens"
+              f"{', answer forced' if forced else ''}; protocol: {src})", flush=True)
+    else:
+        print("[2] backbone alone      : skipped (--no-baseline)")
+
+    # ---- 3. the physics model on the true initial condition
+    t3 = time.time()
+    params = mx.array([[item["a"], math.sin(item["phi"]), math.cos(item["phi"])]], dtype=mx.float32)
+    ic = build_ic_mlx(params)                                   # the bridges' IC parameterization
+    u_field = np.array(fno(ic), dtype=np.float64)[0]            # u(x, t=0.5) on the 128-grid
+    u_fno = fourier_interp(u_field, item["x0"])
+    t_fno = time.time() - t3
+    u_true = fourier_interp(solve(initial_condition(item["a"], item["phi"])), item["x0"])
+    if args.verbose:
+        ic_np = initial_condition(item["a"], item["phi"])
+        print(f"    (IC parameterization check: max|build_ic - initial_condition| = "
+              f"{float(np.abs(np.array(ic)[0] - ic_np).max()):.2e})")
+    print(f"[3] physics model (FNO) : u({item['x0']}) = {u_fno:+.4f}   ({t_fno*1e3:.0f} ms; "
+          f"the reference the coupled answer should match to +-{TOL})")
+    print(f"    spectral solver     : u({item['x0']}) = {u_true:+.4f}   (ground truth)")
+
+    # ---- verdict
+    def mark(v):
+        return "n/a" if v is None else ("match" if abs(v - u_fno) <= TOL else "off")
+    if psi_val is None:
+        verdict = "PsiLM n/a (no number in the reply)"
+    else:
+        verdict = (f"PsiLM {mark(psi_val)} (|{psi_val}-{u_fno:.2f}| "
+                   f"{'<=' if abs(psi_val-u_fno) <= TOL else '>'} {TOL})")
+    print("\n" + verdict + ("" if args.no_baseline else f"; backbone alone {mark(base_val)}"))
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,18 @@
+# PsiLM on Gemma 4 12B -- Apple Silicon (MLX). Verified 2026-09-06 on macOS / Python 3.11
+# with: mlx 0.32.2, mlx-lm 0.31.3, transformers 5.16.1, torch 2.13.0, huggingface_hub 1.29.0.
+#
+#   pip install -r requirements.txt
+#   python psilm_infer.py
+#
+mlx>=0.32
+mlx-lm==0.31.3          # exact: the Gemma 4 loader (psilm.mlx.gemma_loader) goes through
+                        # mlx_lm.utils internals; other versions are untested
+transformers>=4.51      # chat template only (apply_chat_template); the model runs in MLX
+torch>=2.4              # imported by psilm.stage2.qa (batching helpers); CPU wheel is fine
+huggingface_hub>=0.30   # downloads the backbone from mlx-community
+numpy>=1.26
+safetensors>=0.4
+
+# The PsiLM package itself (bridges, staged forward, Gemma loader, FNO, QA builder).
+# Alternative to this line: clone the repository and set PSILM_REPO=/path/to/PsiLM.
+psilm @ git+https://github.com/ryoji-info/PsiLM