Upload jev_toy/serve.py with huggingface_hub

jev_toy/serve.py

@@ -0,0 +1,133 @@
+"""
+jev_toy/serve.py
+
+A mirror of the System One request/response API, backed by our toy model.
+Shows the proven interface in action:
+
+    POST: { state: "...", questions: { <name>: {type, instructions, ...} } }
+    ->     { <name>: { <typed result with probabilities> } }
+
+All questions are answered from ONE encoder pass over the state (parallel).
+
+Usage:
+  source .venv/bin/activate
+  python -m jev_toy.serve --ckpt checkpoints/model.pt
+"""
+
+from __future__ import annotations
+
+import argparse
+
+import torch
+import torch.nn.functional as F
+
+from jev_toy.data import tokenize
+from jev_toy.model import SystemOneConfig, SystemOneModel
+
+
+def build_stoi(vocab_dict):
+    return vocab_dict
+
+
+class ToyServer:
+    def __init__(self, ckpt, device="cpu"):
+        self.device = device
+        ck = torch.load(ckpt, map_location=device)
+        self.cfg = SystemOneConfig(**ck["config"])
+        self.stoi = ck["vocab"]
+        # __oov__ key stores the oov index; default 1
+        self.oov = ck["vocab"].get("__oov__", 1)
+        self.pad = self.cfg.pad_token_id
+        self.model = SystemOneModel(self.cfg).to(device)
+        self.model.load_state_dict(ck["state_dict"])
+        self.model.eval()
+
+    def _enc(self, text, keep=None):
+        ids, mask = tokenize(text, self.stoi, self.cfg.max_seq_len, self.oov, self.pad)
+        return (
+            torch.tensor([ids], device=self.device),
+            torch.tensor([mask], device=self.device),
+        )
+
+    @torch.no_grad()
+    def answer(self, state: str, questions: dict):
+        """questions: {name: {type, instructions, options?}}. Returns {name: decision}."""
+        s_ids, s_mask = self._enc(state)
+        h_state = self.model.encode_state(s_ids, s_mask)  # ONE pass over state
+        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"]
+            if t == "noul":
+                p = torch.sigmoid(logits["noul"][0]).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]
+                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)
+                u = 1.0 / k
+                conf = float((probs.max() - u) / (1 - u))
+                results[name] = {
+                    "type": "choice", "distribution": dist,
+                    "argmax": str(opts[int(probs.argmax())]),
+                    "confidence": round(conf, 4),
+                }
+            elif t == "score":
+                lo, hi = self.cfg.score_range
+                p = (hi - lo) * torch.sigmoid(logits["score"][0]) + lo
+                results[name] = {"type": "score", "score": round(float(p), 3)}
+            else:
+                raise ValueError(t)
+        return results
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--ckpt", default="checkpoints/model.pt")
+    args = ap.parse_args()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    server = ToyServer(args.ckpt, device)
+
+    # --- demo 1: support-ticket-style multi-question (the Jev flagship use case) ---
+    demo_state = "I've been trying to connect my Stripe account for 3 days, it keeps failing with a 403. I'm losing sales and my manager is anxious. Please help ASAP."
+    demo_qs = {
+        "is_urgent": {"type": "noul", "instructions": "The message conveys urgency or time-sensitivity."},
+        "department": {
+            "type": "choice",
+            "instructions": "Which team should handle this?",
+            "options": ["billing", "technical", "sales"],
+        },
+        "frustration": {
+            "type": "score",
+            "instructions": "How frustrated does the customer appear? (0=civil, 1=angry)",
+        },
+    }
+    print("=== demo: one state, three parallel typed questions ===")
+    out = server.answer(demo_state, demo_qs)
+    for k, v in out.items():
+        print(f"  {k}: {v}")
+
+    # --- demo 2: routing a news headline ---
+    print("\n=== demo: route a news article ===")
+    news = "News article: NVIDIA stock hits all-time high after beating earnings estimates, analysts raise price targets."
+    ro = server.answer(news, {
+        "topic": {
+            "type": "choice", "instructions": "Which topic does this article belong to?",
+            "options": ["World", "Sports", "Business", "Sci/Tech"],
+        }
+    })
+    print("  topic:", ro["topic"])
+
+
+if __name__ == "__main__":
+    main()