Upload jev_toy/serve.py with huggingface_hub

jev_toy/serve.py

@@ -51,29 +51,53 @@ class ToyServer:
 
     @torch.no_grad()
     def answer(self, state: str, questions: dict):
-        """questions: {name: {type, instructions, options?}}. Returns {name: decision}."""
+        """
+        questions: {name: {type, instructions, options?}}. Returns {name: decision}.
+
+        PARALLEL PATH (mirrors Jev's proven claim): the state is encoded ONCE
+        (one encoder pass), and ALL questions are encoded in a SINGLE shared
+        encoder pass, then fused with the state vector and sent to the typed
+        heads. Nothing is generated; all answers come from one forward pass on
+        the state + one forward pass on the question batch.
+        """
         s_ids, s_mask = self._enc(state)
-        h_state = self.model.encode_state(s_ids, s_mask)  # ONE pass over state
+        h_state = self.model.encode_state(s_ids, s_mask)      # ONE pass over state
+
+        # encode ALL questions in ONE batched shared-encoder pass (parallel)
+        names = list(questions.keys())
+        q_ids_all, q_mask_all, types_all = [], [], []
+        for name in names:
+            ids, mask = self._enc(questions[name].get("instructions", name))
+            q_ids_all.append(ids[0]); q_mask_all.append(mask[0])
+            types_all.append(questions[name]["type"])
+        q_ids = torch.stack(q_ids_all).to(self.device)         # [N, T]
+        q_mask = torch.stack(q_mask_all).to(self.device)
+        h_q = self.model.encode_questions(q_ids, q_mask)      # ONE pass over all questions
+
+        # fuse state with every question row and apply typed heads
+        h = F.gelu(self.model.merge(torch.cat([h_state.expand(q_ids.shape[0], -1), h_q], dim=-1)))
+        logits = {
+            "noul": self.model.noul_head(h).squeeze(-1),
+            "score": self.model.score_head(h).squeeze(-1),
+            "choice": self.model.choice_head(h),
+        }
+
         results = {}
-        for name, q in questions.items():
-            q_ids, q_mask = self._enc(q.get("instructions", name))  # type: ignore
-            # h_state is [1, D]; one question => one row, so pass it directly.
-            logits, _ = self.model.answer(h_state, q_ids, q_mask, q["type"])
-            t = q["type"]
+        for i, name in enumerate(names):
+            t = types_all[i]
             if t == "noul":
-                p = torch.sigmoid(logits["noul"][0]).item()
+                p = torch.sigmoid(logits["noul"][i]).item()
                 results[name] = {"type": "noul", "noul": round(p, 4), "is_true": p >= 0.5}
             elif t == "choice":
-                opts = q.get("options") or []
-                probs = logits["choice"][0]
+                opts = questions[name].get("options") or []
+                probs = logits["choice"][i]
                 k = len(opts)
                 if k == 0:
                     k = probs.shape[0]
                     opts = [f"option_{i}" for i in range(k)]
                 probs = probs[:k]
-                probs = F.softmax(probs / probs.sum().clamp_min(1e-9), dim=-1)  # renormalize over given options
-                dist = {str(opts[i]): float(f"{probs[i].item():.4f}") for i in range(k)}
-                # confidence = margin above uniform (community formula for choice confidence)
+                probs = F.softmax(probs / probs.sum().clamp_min(1e-9), dim=-1)
+                dist = {str(opts[j]): float(f"{probs[j].item():.4f}") for j in range(k)}
                 u = 1.0 / k
                 conf = float((probs.max() - u) / (1 - u))
                 results[name] = {
@@ -83,16 +107,47 @@ class ToyServer:
                 }
             elif t == "score":
                 lo, hi = self.cfg.score_range
-                p = (hi - lo) * torch.sigmoid(logits["score"][0]) + lo
+                p = (hi - lo) * torch.sigmoid(logits["score"][i]) + lo
                 results[name] = {"type": "score", "score": round(float(p), 3)}
             else:
                 raise ValueError(t)
         return results
 
 
+def benchmark(server, state, questions, iters=20):
+    """Empirically show the parallel claim: encoding the state is done ONCE,
+    regardless of how many questions we ask. We time the full answer() for
+    increasing question counts and show marginal cost per extra question.
+    """
+    import time
+    # build up to 12 questions (reuse the 3 real ones + synthetic nouls)
+    qs = list(questions.items())
+    while len(qs) < 12:
+        qs.append((f"q{len(qs)}", {
+            "type": "noul",
+            "instructions": f"Is the following claim true? Additional check {len(qs)}.",
+        }))
+    rows = []
+    for n in [1, 3, 6, 12]:
+        subset = dict(qs[:n])
+        server.answer(state, subset)  # warmup
+        t0 = time.perf_counter()
+        for _ in range(iters):
+            server.answer(state, subset)
+        dt = (time.perf_counter() - t0) / iters * 1000
+        rows.append((n, dt))
+    print("\n=== benchmark: parallel prediction ===")
+    print("questions | full answer ms (CPU, single request)")
+    for n, ms in rows:
+        print(f"  {n:>8} | {ms:>22.2f}")
+    print("\nKey: adding questions grows roughly with the question-batch size,")
+    print("NOT with re-reading the state. The state is encoded exactly once.")
+
+
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--ckpt", default="checkpoints/model.pt")
+    ap.add_argument("--bench", action="store_true", help="run parallel-encoding benchmark")
     args = ap.parse_args()
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -128,6 +183,9 @@ def main():
     })
     print("  topic:", ro["topic"])
 
+    if args.bench:
+        benchmark(server, demo_state, demo_qs)
+
 
 if __name__ == "__main__":
     main()