Share Qwen3.5 prefix cache across NLI hypotheses

README.md

@@ -1,90 +1,103 @@
----
-license: mit
-base_model: Qwen/Qwen3.5-4B
-pipeline_tag: text-classification
-library_name: transformers
-tags:
-- nli
-- cross-encoder
-- qwen3.5
-- reranker
-- text-classification
-- image-text-to-text
-language:
-- en
----
-
-
-# openjev — Qwen3.5 trained as jev model
-
-<video controls src="https://huggingface.co/AlexWortega/openjev/resolve/main/videos/v2/doom_vision_v2.mp4" width="720"></video>
-
-<video controls src="https://huggingface.co/AlexWortega/openjev/resolve/main/videos/minecraft_chain.mp4" width="720"></video>
-
-**openjev** is Qwen3.5 turned into a *jev* model: a single cross-encoder that reads a premise and a hypothesis and
-answers with entailment, contradiction or neutral. That one primitive is enough to rerank answers, grade them against a
-reference, guard content, and play games in real time: hand it the game state and a few statements about it, and the
-argmax entailment is the move. Nothing is trained per task.
-
-## openjev-4B v2: text, images and agents
-
-The new 4B checkpoint (`qwen3.5-4b-nli-v2/`) reads images as well as text and was trained on a much larger and harder
-mixture. It is strictly zero-shot on everything shown here.
-
-* Doom straight from the pixels (first video): **10.4 kills** per episode, twice the v1 model (5.2); random play gets 1.
-* Crafts an **iron pickaxe from nothing in real Minecraft** (second video): 11 milestones in ~22 decisions, driven by a
-  backward-chaining scaffold where the jev model only checks statements about the inventory and the world.
-* Much stronger on adversarial NLI (ANLI r3 0.42 → 0.63, WANLI 0.63 → 0.77) and on image claims (0.52 → 0.84),
-  better reranking (ARC-Challenge 0.59 → 0.72, MMLU 0.47 → 0.53), same MNLI (0.91).
-
-Doom from the text state (v2, 11 kills per episode; a perfect-information bot gets 18.8):
-
-<video controls src="https://huggingface.co/AlexWortega/openjev/resolve/main/videos/v2/doom_position_v2.mp4" width="720"></video>
-
-![radar](assets/radar_openjev.png)
-
-Bigger jev: Qwen3.5-35B-A3B (MoE) as the backbone (`qwen3.5-35b-a3b-nli/`). Zero-shot, and with the backbone frozen plus
-a small MLP head on the last-token latent (`mlp_heads_35b/`, one head per task, loadable with `LatentMLPHead.load`):
-
-![radar 35B](assets/radar_openjev_35b.png)
-
-## What's inside
-
-* `qwen3.5-4b-nli-v2/` — **recommended**: the 4B v2 jev checkpoint, text + images.
-* `qwen3.5-4b-nli/` — the original 4B jev checkpoint (text).
-* `qwen3.5-35b-a3b-nli/` — the 35B-A3B MoE jev checkpoint (load with `modeling_qwen35_moe_seqcls.py`).
-* All checkpoints: `Qwen3_5ForSequenceClassification`, 3 labels `contradiction`, `entailment`, `neutral`, last-token
-  pooling, trained with plain cross-entropy over the three classes.
-* `modeling_openjev.py` — `OpenJevCrossEncoder`: `predict`, `rerank`, `grade`, `latents`; `LatentMLPHead` for the per-task heads.
-* `modeling_qwen35_moe_seqcls.py` — `Qwen3_5MoeForSequenceClassification` for the 35B-A3B backbone.
-* `mlp_heads_35b/<task>/` — `head.pt` + `norm.npz` + `meta.json`, the 35B latent + MLP heads behind the second radar.
-* `code/` — everything used here: the trainer and data mixture builder, the evaluation harness, Flappy Bird, Doom
-  (text and pixels), the Minecraft scaffold and bot, the radar.
-* `videos/` — Flappy Bird, Doom and Minecraft replays; `results/` — raw JSON for every run and the full report.
-
-## Use it
-
-```python
-from modeling_openjev import OpenJevCrossEncoder
-jev = OpenJevCrossEncoder("AlexWortega/openjev", subfolder="qwen3.5-4b-nli-v2")
-
-jev.predict([("The bird is 0.05 below the centre of the gap.", "The bird is below the centre of the gap.")])
-# -> [[contradiction, entailment, neutral]] probabilities
-
+---
+license: mit
+base_model: Qwen/Qwen3.5-4B
+pipeline_tag: text-classification
+library_name: transformers
+tags:
+- nli
+- cross-encoder
+- qwen3.5
+- reranker
+- text-classification
+- image-text-to-text
+language:
+- en
+---
+
+
+# openjev — Qwen3.5 trained as jev model
+
+<video controls src="https://huggingface.co/AlexWortega/openjev/resolve/main/videos/v2/doom_vision_v2.mp4" width="720"></video>
+
+<video controls src="https://huggingface.co/AlexWortega/openjev/resolve/main/videos/minecraft_chain.mp4" width="720"></video>
+
+**openjev** is Qwen3.5 turned into a *jev* model: a single cross-encoder that reads a premise and a hypothesis and
+answers with entailment, contradiction or neutral. That one primitive is enough to rerank answers, grade them against a
+reference, guard content, and play games in real time: hand it the game state and a few statements about it, and the
+argmax entailment is the move. Nothing is trained per task.
+
+## openjev-4B v2: text, images and agents
+
+The new 4B checkpoint (`qwen3.5-4b-nli-v2/`) reads images as well as text and was trained on a much larger and harder
+mixture. It is strictly zero-shot on everything shown here.
+
+* Doom straight from the pixels (first video): **10.4 kills** per episode, twice the v1 model (5.2); random play gets 1.
+* Crafts an **iron pickaxe from nothing in real Minecraft** (second video): 11 milestones in ~22 decisions, driven by a
+  backward-chaining scaffold where the jev model only checks statements about the inventory and the world.
+* Much stronger on adversarial NLI (ANLI r3 0.42 → 0.63, WANLI 0.63 → 0.77) and on image claims (0.52 → 0.84),
+  better reranking (ARC-Challenge 0.59 → 0.72, MMLU 0.47 → 0.53), same MNLI (0.91).
+
+Doom from the text state (v2, 11 kills per episode; a perfect-information bot gets 18.8):
+
+<video controls src="https://huggingface.co/AlexWortega/openjev/resolve/main/videos/v2/doom_position_v2.mp4" width="720"></video>
+
+![radar](assets/radar_openjev.png)
+
+Bigger jev: Qwen3.5-35B-A3B (MoE) as the backbone (`qwen3.5-35b-a3b-nli/`). Zero-shot, and with the backbone frozen plus
+a small MLP head on the last-token latent (`mlp_heads_35b/`, one head per task, loadable with `LatentMLPHead.load`):
+
+![radar 35B](assets/radar_openjev_35b.png)
+
+## What's inside
+
+* `qwen3.5-4b-nli-v2/` — **recommended**: the 4B v2 jev checkpoint, text + images.
+* `qwen3.5-4b-nli/` — the original 4B jev checkpoint (text).
+* `qwen3.5-35b-a3b-nli/` — the 35B-A3B MoE jev checkpoint (load with `modeling_qwen35_moe_seqcls.py`).
+* All checkpoints: `Qwen3_5ForSequenceClassification`, 3 labels `contradiction`, `entailment`, `neutral`, last-token
+  pooling, trained with plain cross-entropy over the three classes.
+* `modeling_openjev.py` — `OpenJevCrossEncoder`: `predict`, `predict_hypotheses`, `rerank`, `grade`,
+  `latents`, `latents_hypotheses`; `LatentMLPHead` for the per-task heads.
+* `modeling_qwen35_moe_seqcls.py` — `Qwen3_5MoeForSequenceClassification` for the 35B-A3B backbone.
+* `mlp_heads_35b/<task>/` — `head.pt` + `norm.npz` + `meta.json`, the 35B latent + MLP heads behind the second radar.
+* `code/` — everything used here: the trainer and data mixture builder, the evaluation harness, Flappy Bird, Doom
+  (text and pixels), the Minecraft scaffold and bot, the radar.
+* `videos/` — Flappy Bird, Doom and Minecraft replays; `results/` — raw JSON for every run and the full report.
+
+## Use it
+
+```python
+from modeling_openjev import OpenJevCrossEncoder
+jev = OpenJevCrossEncoder("AlexWortega/openjev", subfolder="qwen3.5-4b-nli-v2")
+
+jev.predict([("The bird is 0.05 below the centre of the gap.", "The bird is below the centre of the gap.")])
+# -> [[contradiction, entailment, neutral]] probabilities
+
 jev.rerank("Which gas do plants absorb during photosynthesis?", ["oxygen", "carbon dioxide", "nitrogen"])
 # -> index of the option with the highest entailment
-```
 
-Or with plain transformers:
-
-```python
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-tok = AutoTokenizer.from_pretrained("AlexWortega/openjev", subfolder="qwen3.5-4b-nli-v2")
-model = AutoModelForSequenceClassification.from_pretrained("AlexWortega/openjev", subfolder="qwen3.5-4b-nli-v2")
-text = model.config.nli_template.format(premise="...", hypothesis="...")
+jev.predict_hypotheses("Which gas do plants absorb during photosynthesis?",
+                       ["The correct answer is: oxygen", "The correct answer is: carbon dioxide",
+                        "The correct answer is: nitrogen"])
+# -> one [contradiction, entailment, neutral] row per hypothesis
 ```
 
-Images go inside the premise as `<|vision_start|><|image_pad|>…<|vision_end|>` with `pixel_values` / `image_grid_thw`
-from the Qwen3.5 image processor; see `code/doom_vision.py` and `code/eval_image_nli.py`.
-
-Reference point: [dleemiller's NLI cross-encoders](https://huggingface.co/blog/dleemiller/nli-xenc-ways-to-use). Licence MIT.
+`predict_hypotheses` and `latents_hypotheses` use the existing pairwise batch for one or two hypotheses. For three or
+more, they compute the common token prefix once, then score every hypothesis in one batched continuation. `rerank`
+uses the same rule. Qwen3.5 has recurrent linear-attention layers, so a 4D packed
+tree mask alone would mix branches; the shared prefix cache is copied into separate batch entries for the suffixes.
+This path is for text inputs. Its batch size and suffix padding use memory proportional to the number and length of
+the hypotheses; split very large option sets into smaller calls.
+
+Or with plain transformers:
+
+```python
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+tok = AutoTokenizer.from_pretrained("AlexWortega/openjev", subfolder="qwen3.5-4b-nli-v2")
+model = AutoModelForSequenceClassification.from_pretrained("AlexWortega/openjev", subfolder="qwen3.5-4b-nli-v2")
+text = model.config.nli_template.format(premise="...", hypothesis="...")
+```
+
+Images go inside the premise as `<|vision_start|><|image_pad|>…<|vision_end|>` with `pixel_values` / `image_grid_thw`
+from the Qwen3.5 image processor; see `code/doom_vision.py` and `code/eval_image_nli.py`.
+
+Reference point: [dleemiller's NLI cross-encoders](https://huggingface.co/blog/dleemiller/nli-xenc-ways-to-use). Licence MIT.

modeling_openjev.py

@@ -1,197 +1,261 @@
-"""openjev: a Qwen3.5 decoder used as an NLI cross-encoder, plus a latent + MLP head trained with soft BCE.
-
-Architecture
-------------
-* `OpenJevCrossEncoder`: `Qwen3_5ForSequenceClassification` — the Qwen3.5 text backbone followed by a linear
-  `score` head over the hidden state of the last non-pad token. Three labels in dleemiller order
-  (0 = contradiction, 1 = entailment, 2 = neutral). Input is one string
-  `"Premise: {premise}\\nHypothesis: {hypothesis}"` (template stored in `config.nli_template`), right-padded.
-  The vision tower that ships with Qwen3.5 checkpoints is kept but unused for text. Loss for the NLI head:
-  cross-entropy over the three classes.
-* `LatentMLPHead`: a small MLP (d -> 512 -> 1, GELU, dropout 0.1) on top of the frozen cross-encoder latent
-  (the pooled hidden state that feeds `score`), one scalar per (question, option) pair. Loss: soft BCE —
-  `BCEWithLogits(logit, y*(1-eps) + (1-y)*eps)` with positives re-weighted by (1-p)/p, p = positive rate.
-  At inference the option with the highest score wins (per-question argmax).
-
-Usage
------
-    from modeling_openjev import OpenJevCrossEncoder, LatentMLPHead
-    ce = OpenJevCrossEncoder("AlexWortega/openjev", subfolder="qwen3.5-4b-nli")
-    ce.predict([("A man is playing a guitar.", "Someone is making music.")])   # -> [[p_con, p_ent, p_neu]]
-    ce.rerank("What is the capital of France?", ["Paris", "Lyon", "Berlin"])   # -> index of the best option
-
-    X = ce.latents(pairs)                       # (n_pairs, d) float32
-    head = LatentMLPHead(X.shape[1]).fit(X, gold, qid)   # gold in {0,1} per pair, qid groups pairs by question
-    head.predict(X)                             # (n_pairs,) logits; argmax within each qid
-"""
-from __future__ import annotations
-
-import json
-import os
-
-import numpy as np
-import torch
-import torch.nn as nn
-
-CON, ENT, NEU = 0, 1, 2
-DEFAULT_TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"
-
-
-class OpenJevCrossEncoder:
-    def __init__(self, path: str, subfolder: str | None = None, device: str | None = None, dtype=torch.bfloat16,
-                 bs: int = 32, max_len: int = 4096):
-        from transformers import AutoModelForSequenceClassification, AutoTokenizer
-
-        kw = {"subfolder": subfolder} if subfolder else {}
-        self.tok = AutoTokenizer.from_pretrained(path, **kw)
-        self.model = AutoModelForSequenceClassification.from_pretrained(path, dtype=dtype, **kw)
-        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.model.to(self.device).eval()
-        self.template = getattr(self.model.config, "nli_template", None) or DEFAULT_TEMPLATE
-        if self.tok.pad_token is None:
-            self.tok.pad_token = self.tok.eos_token
-        self.tok.padding_side = "right"  # the head pools the last non-pad token
-        tc = self.model.config.get_text_config()
-        if tc.pad_token_id is None:
-            tc.pad_token_id = self.tok.pad_token_id
-        self.backbone = getattr(self.model, self.model.base_model_prefix)
-        self.bs, self.max_len = bs, max_len
-
-    def _encode(self, pairs):
-        texts = [self.template.format(premise=p.strip(), hypothesis=h.strip()) for p, h in pairs]
-        enc = self.tok(texts, truncation=True, max_length=self.max_len, padding=True, return_tensors="pt")
-        return {k: v.to(self.device) for k, v in enc.items()}
-
-    @torch.no_grad()
-    def _pooled(self, enc):
-        h = self.backbone(**enc).last_hidden_state
-        last = enc["attention_mask"].sum(1) - 1
-        return h[torch.arange(h.shape[0], device=h.device), last]
-
-    @torch.no_grad()
-    def latents(self, pairs) -> np.ndarray:
-        """Pooled last-token hidden state (the input of the `score` head) for each (premise, hypothesis) pair."""
-        out = []
-        for i in range(0, len(pairs), self.bs):
-            out.append(self._pooled(self._encode(pairs[i:i + self.bs])).float().cpu().numpy())
-        return np.concatenate(out, 0)
-
-    @torch.no_grad()
+"""openjev: a Qwen3.5 decoder used as an NLI cross-encoder, plus a latent + MLP head trained with soft BCE.
+
+Architecture
+------------
+* `OpenJevCrossEncoder`: `Qwen3_5ForSequenceClassification` — the Qwen3.5 text backbone followed by a linear
+  `score` head over the hidden state of the last non-pad token. Three labels in dleemiller order
+  (0 = contradiction, 1 = entailment, 2 = neutral). Input is one string
+  `"Premise: {premise}\\nHypothesis: {hypothesis}"` (template stored in `config.nli_template`), right-padded.
+  The vision tower that ships with Qwen3.5 checkpoints is kept but unused for text. Loss for the NLI head:
+  cross-entropy over the three classes.
+* `LatentMLPHead`: a small MLP (d -> 512 -> 1, GELU, dropout 0.1) on top of the frozen cross-encoder latent
+  (the pooled hidden state that feeds `score`), one scalar per (question, option) pair. Loss: soft BCE —
+  `BCEWithLogits(logit, y*(1-eps) + (1-y)*eps)` with positives re-weighted by (1-p)/p, p = positive rate.
+  At inference the option with the highest score wins (per-question argmax).
+
+Usage
+-----
+    from modeling_openjev import OpenJevCrossEncoder, LatentMLPHead
+    ce = OpenJevCrossEncoder("AlexWortega/openjev", subfolder="qwen3.5-4b-nli")
+    ce.predict([("A man is playing a guitar.", "Someone is making music.")])   # -> [[p_con, p_ent, p_neu]]
+    ce.rerank("What is the capital of France?", ["Paris", "Lyon", "Berlin"])   # -> index of the best option
+
+    X = ce.latents(pairs)                       # (n_pairs, d) float32
+    head = LatentMLPHead(X.shape[1]).fit(X, gold, qid)   # gold in {0,1} per pair, qid groups pairs by question
+    head.predict(X)                             # (n_pairs,) logits; argmax within each qid
+"""
+from __future__ import annotations
+
+import json
+import os
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+CON, ENT, NEU = 0, 1, 2
+DEFAULT_TEMPLATE = "Premise: {premise}\nHypothesis: {hypothesis}"
+
+
+class OpenJevCrossEncoder:
+    def __init__(self, path: str, subfolder: str | None = None, device: str | None = None, dtype=torch.bfloat16,
+                 bs: int = 32, max_len: int = 4096):
+        from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+        kw = {"subfolder": subfolder} if subfolder else {}
+        self.tok = AutoTokenizer.from_pretrained(path, **kw)
+        self.model = AutoModelForSequenceClassification.from_pretrained(path, dtype=dtype, **kw)
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device).eval()
+        self.template = getattr(self.model.config, "nli_template", None) or DEFAULT_TEMPLATE
+        if self.tok.pad_token is None:
+            self.tok.pad_token = self.tok.eos_token
+        self.tok.padding_side = "right"  # the head pools the last non-pad token
+        tc = self.model.config.get_text_config()
+        if tc.pad_token_id is None:
+            tc.pad_token_id = self.tok.pad_token_id
+        self.backbone = getattr(self.model, self.model.base_model_prefix)
+        self.bs, self.max_len = bs, max_len
+
+    def _encode(self, pairs):
+        texts = [self.template.format(premise=p.strip(), hypothesis=h.strip()) for p, h in pairs]
+        enc = self.tok(texts, truncation=True, max_length=self.max_len, padding=True, return_tensors="pt")
+        return {k: v.to(self.device) for k, v in enc.items()}
+
+    @torch.no_grad()
+    def _pooled(self, enc):
+        h = self.backbone(**enc).last_hidden_state
+        last = enc["attention_mask"].sum(1) - 1
+        return h[torch.arange(h.shape[0], device=h.device), last]
+
+    @torch.no_grad()
+    def latents(self, pairs) -> np.ndarray:
+        """Pooled last-token hidden state (the input of the `score` head) for each (premise, hypothesis) pair."""
+        out = []
+        for i in range(0, len(pairs), self.bs):
+            out.append(self._pooled(self._encode(pairs[i:i + self.bs])).float().cpu().numpy())
+        return np.concatenate(out, 0)
+
+    @torch.no_grad()
     def predict(self, pairs) -> np.ndarray:
-        """Softmax probabilities [contradiction, entailment, neutral] per pair."""
-        out = []
-        for i in range(0, len(pairs), self.bs):
-            logits = self.model.score(self._pooled(self._encode(pairs[i:i + self.bs]))).float()
-            out.append(torch.softmax(logits, -1).cpu().numpy())
+        """Softmax probabilities [contradiction, entailment, neutral] per pair."""
+        out = []
+        for i in range(0, len(pairs), self.bs):
+            logits = self.model.score(self._pooled(self._encode(pairs[i:i + self.bs]))).float()
+            out.append(torch.softmax(logits, -1).cpu().numpy())
         return np.concatenate(out, 0)
 
-    def rerank(self, question: str, options, hyp_fmt: str = "The correct answer is: {}") -> int:
-        """Zero-shot multiple choice: option with the highest P(entailment) given the question as premise."""
-        p = self.predict([(question, hyp_fmt.format(o)) for o in options])
-        return int(p[:, ENT].argmax())
-
-    def grade(self, question: str, reference: str, candidate: str) -> str:
-        """Reference-based grading: premise = question + reference, hypothesis = candidate."""
-        p = self.predict([(f"{question}\nReference answer: {reference}", f"Answer: {candidate}")])[0]
-        return ["contradiction", "entailment", "neutral"][int(p.argmax())]
-
-
-class _MLP(nn.Module):
-    def __init__(self, d, hidden=512, p=0.1):
-        super().__init__()
-        self.net = nn.Sequential(nn.Linear(d, hidden), nn.GELU(), nn.Dropout(p), nn.Linear(hidden, 1))
-
-    def forward(self, x):
-        return self.net(x).squeeze(-1)
-
-
-def soft_bce(logits, y, eps, pos_weight):
-    """BCE with soft targets: gold -> 1-eps, others -> eps; positives up-weighted by pos_weight."""
-    target = y * (1 - eps) + (1 - y) * eps
-    w = torch.where(y > 0.5, pos_weight, torch.ones_like(y))
-    return (w * nn.functional.binary_cross_entropy_with_logits(logits, target, reduction="none")).mean()
-
-
-def per_question_acc(scores, qid, gold):
-    """Fraction of questions whose argmax-scored option is the gold one."""
-    order = np.argsort(qid, kind="stable")
-    scores, qid, gold = scores[order], qid[order], gold[order]
-    starts = np.r_[0, np.flatnonzero(np.diff(qid)) + 1, len(qid)]
-    hits = [gold[a:b][scores[a:b].argmax()] == 1 for a, b in zip(starts[:-1], starts[1:])]
-    return float(np.mean(hits))
-
-
-def grouped_split(qid, frac, seed):
-    qs = np.unique(qid)
-    rng = np.random.RandomState(seed)
-    rng.shuffle(qs)
-    hold = set(qs[: max(1, int(len(qs) * frac))].tolist())
-    mask = np.array([q in hold for q in qid])
-    return ~mask, mask
+    @torch.no_grad()
+    def _pooled_hypotheses(self, premise: str, hypotheses) -> torch.Tensor:
+        """Score hypotheses with one shared text prefix and a batched cached continuation.
 
+        Qwen3.5 mixes full and recurrent linear attention. A packed 4D tree mask
+        cannot isolate the recurrent layers, so their state and the full-attention
+        KV cache are branched into the batch after a single prefix prefill.
+        """
+        hypotheses = list(hypotheses)
+        if not hypotheses:
+            raise ValueError("hypotheses must not be empty")
+        pairs = [(premise, h) for h in hypotheses]
+        # A second forward and cache branching are unlikely to pay off for
+        # one or two hypotheses; use the existing batched path instead.
+        if len(pairs) < 3:
+            return self._pooled(self._encode(pairs))
 
-class LatentMLPHead:
-    def __init__(self, d: int, hidden: int = 512, dropout: float = 0.1, eps: float = 0.1, lr: float = 1e-3,
-                 wd: float = 1e-2, bs: int = 512, epochs: int = 60, patience: int = 8, seed: int = 0, device: str | None = None):
-        self.cfg = dict(d=d, hidden=hidden, dropout=dropout, eps=eps, lr=lr, wd=wd, bs=bs, epochs=epochs, patience=patience, seed=seed)
-        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        torch.manual_seed(seed)
-        self.model = _MLP(d, hidden, dropout).to(self.device)
-        self.mu, self.sd = np.zeros((1, d), np.float32), np.ones((1, d), np.float32)
+        # Tokenize complete inputs first: a tokenizer can merge characters at
+        # the premise/hypothesis boundary, so a textual split is not reliable.
+        texts = [self.template.format(premise=p.strip(), hypothesis=h.strip()) for p, h in pairs]
+        sequences = [self.tok(t, truncation=True, max_length=self.max_len)["input_ids"] for t in texts]
+        common = 0
+        for column in zip(*sequences):
+            if len(set(column)) != 1:
+                break
+            common += 1
+        common = min(common, min(map(len, sequences)) - 1)
+        if common < 1:
+            return self._pooled(self._encode(pairs))
 
-    def _t(self, X):
-        return torch.tensor((X - self.mu) / self.sd, dtype=torch.float32, device=self.device)
+        prefix = torch.tensor([sequences[0][:common]], device=self.device)
+        prefix_positions = torch.arange(common, device=self.device).unsqueeze(0)
+        cache = self.backbone(input_ids=prefix, position_ids=prefix_positions, use_cache=True).past_key_values
+        if cache is None:
+            raise RuntimeError("Qwen3.5 did not return a cache for shared-prefix inference")
+        cache.reorder_cache(torch.zeros(len(sequences), dtype=torch.long, device=self.device))
 
-    def fit(self, X, gold, qid, val_frac: float = 0.1):
-        """X: (n_pairs, d); gold: 1 for the correct option of a question, 0 otherwise; qid: question id per pair.
-        Early stopping on per-question accuracy of a grouped hold-out; restores the best weights."""
-        X, gold, qid = np.asarray(X, np.float32), np.asarray(gold, np.float32), np.asarray(qid)
-        tr, va = grouped_split(qid, val_frac, self.cfg["seed"])
-        self.mu, self.sd = X[tr].mean(0, keepdims=True), X[tr].std(0, keepdims=True) + 1e-6
-        Xtr, ytr, Xva = self._t(X[tr]), torch.tensor(gold[tr], device=self.device), self._t(X[va])
-        p = float(gold[tr].mean())
-        pos_weight = torch.tensor((1 - p) / max(p, 1e-6), device=self.device)
-        opt = torch.optim.AdamW(self.model.parameters(), lr=self.cfg["lr"], weight_decay=self.cfg["wd"])
-        best, best_state, bad = -1.0, None, 0
-        for ep in range(self.cfg["epochs"]):
-            self.model.train()
-            perm = torch.randperm(len(Xtr), device=self.device)
-            for s in range(0, len(Xtr), self.cfg["bs"]):
-                idx = perm[s:s + self.cfg["bs"]]
-                loss = soft_bce(self.model(Xtr[idx]), ytr[idx], self.cfg["eps"], pos_weight)
-                opt.zero_grad(); loss.backward(); opt.step()
-            self.model.eval()
-            with torch.no_grad():
-                acc = per_question_acc(self.model(Xva).cpu().numpy(), qid[va], gold[va])
-            if acc > best:
-                best, bad, best_state = acc, 0, {k: v.clone() for k, v in self.model.state_dict().items()}
-            else:
-                bad += 1
-                if bad >= self.cfg["patience"]:
-                    break
-        self.model.load_state_dict(best_state)
-        self.model.eval()
-        self.val_acc = best
-        return self
+        suffixes = [s[common:] for s in sequences]
+        lengths = torch.tensor([len(s) for s in suffixes], device=self.device)
+        width = int(lengths.max())
+        input_ids = torch.full((len(suffixes), width), self.tok.pad_token_id,
+                               dtype=torch.long, device=self.device)
+        for row, suffix in enumerate(suffixes):
+            input_ids[row, :len(suffix)] = torch.tensor(suffix, device=self.device)
+        offsets = torch.arange(width, device=self.device).unsqueeze(0)
+        valid = offsets < lengths.unsqueeze(1)
+        attention_mask = torch.cat((torch.ones((len(suffixes), common), device=self.device,
+                                               dtype=torch.long), valid.long()), dim=1)
+        position_ids = (common + offsets).expand(len(suffixes), -1)
+        hidden = self.backbone(input_ids=input_ids, attention_mask=attention_mask,
+                               position_ids=position_ids, past_key_values=cache,
+                               use_cache=True).last_hidden_state
+        return hidden[torch.arange(len(suffixes), device=self.device), lengths - 1]
 
     @torch.no_grad()
-    def predict(self, X) -> np.ndarray:
-        """One logit per pair; pick the argmax within each question (sigmoid gives a per-option probability)."""
-        self.model.eval()
-        return self.model(self._t(np.asarray(X, np.float32))).cpu().numpy()
+    def latents_hypotheses(self, premise: str, hypotheses) -> np.ndarray:
+        """Pooled last-token latents for hypotheses sharing one premise."""
+        return self._pooled_hypotheses(premise, hypotheses).float().cpu().numpy()
 
-    def save(self, path: str):
-        os.makedirs(path, exist_ok=True)
-        torch.save(self.model.state_dict(), os.path.join(path, "head.pt"))
-        np.savez(os.path.join(path, "norm.npz"), mu=self.mu, sd=self.sd)
-        json.dump(self.cfg, open(os.path.join(path, "config.json"), "w"), indent=2)
-
-    @classmethod
-    def load(cls, path: str, device: str | None = None) -> "LatentMLPHead":
-        cfg = json.load(open(os.path.join(path, "config.json")))
-        head = cls(device=device, **cfg)
-        head.model.load_state_dict(torch.load(os.path.join(path, "head.pt"), map_location=head.device))
-        z = np.load(os.path.join(path, "norm.npz"))
-        head.mu, head.sd = z["mu"], z["sd"]
-        head.model.eval()
-        return head
+    @torch.no_grad()
+    def predict_hypotheses(self, premise: str, hypotheses) -> np.ndarray:
+        """NLI probabilities for one premise and many hypotheses, sharing the prefix."""
+        return torch.softmax(self.model.score(self._pooled_hypotheses(premise, hypotheses)).float(), -1).cpu().numpy()
+
+    def rerank(self, question: str, options, hyp_fmt: str = "The correct answer is: {}") -> int:
+        """Zero-shot multiple choice: option with the highest P(entailment) given the question as premise."""
+        p = self.predict_hypotheses(question, [hyp_fmt.format(o) for o in options])
+        return int(p[:, ENT].argmax())
+
+    def grade(self, question: str, reference: str, candidate: str) -> str:
+        """Reference-based grading: premise = question + reference, hypothesis = candidate."""
+        p = self.predict([(f"{question}\nReference answer: {reference}", f"Answer: {candidate}")])[0]
+        return ["contradiction", "entailment", "neutral"][int(p.argmax())]
+
+
+class _MLP(nn.Module):
+    def __init__(self, d, hidden=512, p=0.1):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(d, hidden), nn.GELU(), nn.Dropout(p), nn.Linear(hidden, 1))
+
+    def forward(self, x):
+        return self.net(x).squeeze(-1)
+
+
+def soft_bce(logits, y, eps, pos_weight):
+    """BCE with soft targets: gold -> 1-eps, others -> eps; positives up-weighted by pos_weight."""
+    target = y * (1 - eps) + (1 - y) * eps
+    w = torch.where(y > 0.5, pos_weight, torch.ones_like(y))
+    return (w * nn.functional.binary_cross_entropy_with_logits(logits, target, reduction="none")).mean()
+
+
+def per_question_acc(scores, qid, gold):
+    """Fraction of questions whose argmax-scored option is the gold one."""
+    order = np.argsort(qid, kind="stable")
+    scores, qid, gold = scores[order], qid[order], gold[order]
+    starts = np.r_[0, np.flatnonzero(np.diff(qid)) + 1, len(qid)]
+    hits = [gold[a:b][scores[a:b].argmax()] == 1 for a, b in zip(starts[:-1], starts[1:])]
+    return float(np.mean(hits))
+
+
+def grouped_split(qid, frac, seed):
+    qs = np.unique(qid)
+    rng = np.random.RandomState(seed)
+    rng.shuffle(qs)
+    hold = set(qs[: max(1, int(len(qs) * frac))].tolist())
+    mask = np.array([q in hold for q in qid])
+    return ~mask, mask
+
+
+class LatentMLPHead:
+    def __init__(self, d: int, hidden: int = 512, dropout: float = 0.1, eps: float = 0.1, lr: float = 1e-3,
+                 wd: float = 1e-2, bs: int = 512, epochs: int = 60, patience: int = 8, seed: int = 0, device: str | None = None):
+        self.cfg = dict(d=d, hidden=hidden, dropout=dropout, eps=eps, lr=lr, wd=wd, bs=bs, epochs=epochs, patience=patience, seed=seed)
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        torch.manual_seed(seed)
+        self.model = _MLP(d, hidden, dropout).to(self.device)
+        self.mu, self.sd = np.zeros((1, d), np.float32), np.ones((1, d), np.float32)
+
+    def _t(self, X):
+        return torch.tensor((X - self.mu) / self.sd, dtype=torch.float32, device=self.device)
+
+    def fit(self, X, gold, qid, val_frac: float = 0.1):
+        """X: (n_pairs, d); gold: 1 for the correct option of a question, 0 otherwise; qid: question id per pair.
+        Early stopping on per-question accuracy of a grouped hold-out; restores the best weights."""
+        X, gold, qid = np.asarray(X, np.float32), np.asarray(gold, np.float32), np.asarray(qid)
+        tr, va = grouped_split(qid, val_frac, self.cfg["seed"])
+        self.mu, self.sd = X[tr].mean(0, keepdims=True), X[tr].std(0, keepdims=True) + 1e-6
+        Xtr, ytr, Xva = self._t(X[tr]), torch.tensor(gold[tr], device=self.device), self._t(X[va])
+        p = float(gold[tr].mean())
+        pos_weight = torch.tensor((1 - p) / max(p, 1e-6), device=self.device)
+        opt = torch.optim.AdamW(self.model.parameters(), lr=self.cfg["lr"], weight_decay=self.cfg["wd"])
+        best, best_state, bad = -1.0, None, 0
+        for ep in range(self.cfg["epochs"]):
+            self.model.train()
+            perm = torch.randperm(len(Xtr), device=self.device)
+            for s in range(0, len(Xtr), self.cfg["bs"]):
+                idx = perm[s:s + self.cfg["bs"]]
+                loss = soft_bce(self.model(Xtr[idx]), ytr[idx], self.cfg["eps"], pos_weight)
+                opt.zero_grad(); loss.backward(); opt.step()
+            self.model.eval()
+            with torch.no_grad():
+                acc = per_question_acc(self.model(Xva).cpu().numpy(), qid[va], gold[va])
+            if acc > best:
+                best, bad, best_state = acc, 0, {k: v.clone() for k, v in self.model.state_dict().items()}
+            else:
+                bad += 1
+                if bad >= self.cfg["patience"]:
+                    break
+        self.model.load_state_dict(best_state)
+        self.model.eval()
+        self.val_acc = best
+        return self
+
+    @torch.no_grad()
+    def predict(self, X) -> np.ndarray:
+        """One logit per pair; pick the argmax within each question (sigmoid gives a per-option probability)."""
+        self.model.eval()
+        return self.model(self._t(np.asarray(X, np.float32))).cpu().numpy()
+
+    def save(self, path: str):
+        os.makedirs(path, exist_ok=True)
+        torch.save(self.model.state_dict(), os.path.join(path, "head.pt"))
+        np.savez(os.path.join(path, "norm.npz"), mu=self.mu, sd=self.sd)
+        json.dump(self.cfg, open(os.path.join(path, "config.json"), "w"), indent=2)
+
+    @classmethod
+    def load(cls, path: str, device: str | None = None) -> "LatentMLPHead":
+        cfg = json.load(open(os.path.join(path, "config.json")))
+        head = cls(device=device, **cfg)
+        head.model.load_state_dict(torch.load(os.path.join(path, "head.pt"), map_location=head.device))
+        z = np.load(os.path.join(path, "norm.npz"))
+        head.mu, head.sd = z["mu"], z["sd"]
+        head.model.eval()
+        return head

test_shared_prefix.py

@@ -0,0 +1,83 @@
+"""Small CPU checks for shared-prefix inference; no checkpoint download needed."""
+import unittest
+from unittest.mock import patch
+
+import numpy as np
+import torch
+from transformers import (Qwen3_5Config, Qwen3_5ForSequenceClassification,
+                          Qwen3_5TextConfig, Qwen3_5VisionConfig)
+
+from modeling_openjev import OpenJevCrossEncoder
+
+
+class CharacterTokenizer:
+    pad_token_id = 0
+
+    def __call__(self, texts, truncation=True, max_length=4096, padding=False, return_tensors=None):
+        if isinstance(texts, str):
+            return {"input_ids": [ord(c) + 1 for c in texts][:max_length]}
+        rows = [[ord(c) + 1 for c in text][:max_length] for text in texts]
+        width = max(map(len, rows))
+        ids = [row + [0] * (width - len(row)) for row in rows]
+        mask = [[1] * len(row) + [0] * (width - len(row)) for row in rows]
+        return {"input_ids": torch.tensor(ids), "attention_mask": torch.tensor(mask)}
+
+
+class SharedPrefixTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        torch.manual_seed(7)
+        config = Qwen3_5TextConfig(
+            vocab_size=256, hidden_size=32, intermediate_size=64,
+            num_hidden_layers=4, num_attention_heads=4, num_key_value_heads=2,
+            head_dim=8, linear_num_key_heads=2, linear_num_value_heads=4,
+            linear_key_head_dim=8, linear_value_head_dim=8,
+            layer_types=["linear_attention", "full_attention"] * 2,
+            rope_parameters={"rope_type": "default", "rope_theta": 10000,
+                             "partial_rotary_factor": 0.25, "mrope_section": [1, 0, 0]},
+            pad_token_id=0, num_labels=3, use_cache=False,
+        )
+        vision = Qwen3_5VisionConfig(depth=1, hidden_size=32, intermediate_size=64,
+                                    num_heads=4, out_hidden_size=32)
+        model = Qwen3_5ForSequenceClassification(
+            Qwen3_5Config(text_config=config, vision_config=vision, num_labels=3,
+                           pad_token_id=0)).eval()
+        cls.jev = OpenJevCrossEncoder.__new__(OpenJevCrossEncoder)
+        cls.jev.model = model
+        cls.jev.backbone = model.model
+        cls.jev.tok = CharacterTokenizer()
+        cls.jev.template = "Premise: {premise}\nHypothesis: {hypothesis}"
+        cls.jev.device = "cpu"
+        cls.jev.bs = 32
+        cls.jev.max_len = 128
+
+    def test_shared_matches_independent_and_isolates_branches(self):
+        premise = "A cat sleeps."
+        hypotheses = ["A cat rests.", "The dog runs.", "A cat."]
+        expected = self.jev.latents([(premise, h) for h in hypotheses])
+        actual = self.jev.latents_hypotheses(premise, hypotheses)
+        np.testing.assert_allclose(actual, expected, atol=2e-4, rtol=2e-4)
+        expected_probs = self.jev.predict([(premise, h) for h in hypotheses])
+        actual_probs = self.jev.predict_hypotheses(premise, hypotheses)
+        np.testing.assert_allclose(actual_probs, expected_probs, atol=2e-4, rtol=2e-4)
+
+        changed = hypotheses.copy()
+        changed[1] = "Birds fly around today."
+        other = self.jev.latents_hypotheses(premise, changed)
+        np.testing.assert_allclose(other[[0, 2]], actual[[0, 2]], atol=2e-4, rtol=2e-4)
+
+    def test_small_option_sets_use_pairwise_batch(self):
+        for hypotheses in (["B"], ["B", "C"]):
+            with self.subTest(count=len(hypotheses)):
+                with patch.object(self.jev, "_pooled", wraps=self.jev._pooled) as pooled:
+                    actual = self.jev.predict_hypotheses("A", hypotheses)
+                    self.assertEqual(actual.shape, (len(hypotheses), 3))
+                    pooled.assert_called_once()
+
+    def test_empty(self):
+        with self.assertRaises(ValueError):
+            self.jev.predict_hypotheses("A", [])
+
+
+if __name__ == "__main__":
+    unittest.main()