Cleanup init commit

scripts/ws_lens_smoke.py

@@ -0,0 +1,60 @@
+import asyncio
+import json
+import os
+import sys
+
+import websockets
+
+PORT = os.environ.get("JWASH_PORT", "8381")
+
+PROMPT = (
+    sys.argv[1]
+    if len(sys.argv) > 1
+    else "Fact: The currency used in the country shaped like a boot is what? Answer in one word."
+)
+MAX_TOKENS = int(sys.argv[2]) if len(sys.argv) > 2 else 80
+
+
+async def run_chat(ws, use_lens, max_tokens=MAX_TOKENS):
+    await ws.send(
+        json.dumps(
+            {
+                "type": "chat",
+                "messages": [{"role": "user", "content": PROMPT}],
+                "sampling": {"temperature": 0.7, "max_tokens": max_tokens},
+                "lens": use_lens,
+            }
+        )
+    )
+    reading = thinking = 0
+    sample_frame = None
+    text = ""
+    while True:
+        frame = json.loads(await ws.recv())
+        if frame["type"] == "frame":
+            if frame["phase"] == "reading":
+                reading += 1
+            else:
+                thinking += 1
+                sample_frame = frame
+        elif frame["type"] == "done":
+            return frame, reading, thinking, sample_frame
+        elif frame["type"] == "error":
+            print("[error]", frame["message"])
+            sys.exit(1)
+
+
+async def main():
+    async with websockets.connect(f"ws://127.0.0.1:{PORT}/ws", max_size=None) as ws:
+        done, r, t, sample = await run_chat(ws, True)
+        print(f"with lens: {done['stats']}  reading frames={r} thinking={t}")
+        print(f"reply: {done['text'][:120]!r}")
+        if sample:
+            layer, d = sorted(sample["layers"].items(), key=lambda kv: int(kv[0]))[len(sample["layers"]) // 2]
+            print(f"thinking frame pos={sample['pos']} tok={sample['tok']!r} L{layer}:",
+                  [(s.strip(), round(p, 3), rk) for s, p, rk in zip(d["m_strs"][:5], d["m_p"][:5], d["m_rank"][:5])])
+        done2, _, _, _ = await run_chat(ws, False)
+        print(f"without lens: {done2['stats']}")
+
+
+asyncio.run(main())