Engine: streaming auto and chat

dmm/engine.py

@@ -154,9 +154,13 @@ class Engine:
     # Generation --------------------------------------------------------------------------
 
     @torch.no_grad()
-    def _generate(self, ids: list[int], max_new_tokens: int, cache: StateCache | None = None):
+    def _generate(
+        self, ids: list[int], max_new_tokens: int, cache: StateCache | None = None, streamer=None
+    ):
         x = torch.tensor([ids], device=self.model.device)
         kw = {"past_key_values": cache.kv} if cache is not None else {}
+        if streamer is not None:
+            kw["streamer"] = streamer
         with lora.disabled(self.model):
             out = self.model.generate(
                 x,
@@ -168,6 +172,68 @@ class Engine:
             )
         return out[0, len(ids) :].tolist()
 
+    def _stream(self, ids: list[int], max_new_tokens: int, cache: StateCache | None = None):
+        """Generate as _generate does, yielding {"event": "text"} pieces as they decode.
+
+        Returns the generated token ids (the generator's return value). Generation runs in a
+        thread feeding a streamer; the thread is always joined, even if the caller stops early,
+        so the model and the state cache are never touched by two requests at once.
+        """
+        from transformers import TextIteratorStreamer
+
+        streamer = TextIteratorStreamer(self.tok, skip_prompt=True, skip_special_tokens=True)
+        result: dict = {}
+
+        def run():
+            try:
+                result["tokens"] = self._generate(ids, max_new_tokens, cache, streamer)
+            except BaseException as e:  # surface in the caller's thread
+                result["error"] = e
+                streamer.end()
+
+        t = threading.Thread(target=run, daemon=True)
+        t.start()
+        try:
+            for piece in streamer:
+                if piece:
+                    yield {"event": "text", "text": piece}
+        finally:
+            t.join()
+        if "error" in result:
+            raise result["error"]
+        return result["tokens"]
+
+    def chat_stream(self, messages: list[dict], max_new_tokens: int = 512):
+        """chat(), streamed: {"event": "text"} pieces, then {"event": "done", "tokens": [...]}."""
+        ids = self.tok.apply_chat_template(
+            messages, add_generation_prompt=True, tokenize=True, return_dict=True
+        )["input_ids"]
+        with self.lock:
+            tokens = yield from self._stream(ids, max_new_tokens)
+        yield {"event": "done", "tokens": tokens}
+
+    def auto_stream(self, state: str, name: str, question: dict, max_new_tokens: int = 512):
+        """auto(), streamed. Yields {"event": "decision", ...} as soon as the quick decision and
+        the router's call are known, then {"event": "text"} reasoning pieces when escalated, then
+        {"event": "done", ...} carrying exactly what auto() returns."""
+        item = to_item(state, name, question)
+        with self.lock:
+            cache = StateCache(self.model, self.tpl.state_ids(item))
+            r = self._decide_on(cache, item)
+            out = {
+                "decision": self._answer(item, r),
+                "escalated": r["escalate"],
+                "state_prefills": 1,
+            }
+            yield {"event": "decision", **out}
+            if r["escalate"]:
+                ids = cache.ids + self.tpl.encode(chat_text(item))
+                tokens = yield from self._stream(ids, max_new_tokens, cache)
+                out["generation_tokens"] = tokens
+                out["generation"] = self.tok.decode(tokens, skip_special_tokens=True)
+                cache.kv.crop(len(cache))
+        yield {"event": "done", **out}
+
     def chat(self, messages: list[dict], max_new_tokens: int = 512) -> list[int]:
         ids = self.tok.apply_chat_template(
             messages, add_generation_prompt=True, tokenize=True, return_dict=True