first tests

session.py

@@ -0,0 +1,139 @@
+"""A conversation held in the model's KV cache, with audio prefilled while the user is still speaking.
+
+The cache always ends right after a complete chat-template segment, so each step only appends:
+
+    begin_turn():  [<bos>system…<|im_end|>\\n]  <|im_start|>user\\n
+    feed_audio():  audio tokens, encoded and prefilled as the mic chunks arrive
+    generate():    <|im_end|>\\n<|im_start|>assistant\\n<|mixed_start|>  → generated text + speech frames
+
+The model's own generated tokens (including its speech frames) stay in the cache, so the next turn starts from
+there instead of re-encoding the whole conversation.
+"""
+
+import copy
+import time
+
+import numpy as np
+import soxr
+import torch
+
+MAX_CONTEXT = 16_000  # tokens; past this the conversation is reset at the next turn
+
+
+class StreamingSession:
+    def __init__(self, engine):
+        self.engine = engine
+        self.model = engine.model
+        self.tok = engine.tokenizer
+        self.fe = engine.processor.feature_extractor
+        self.im_end = self.tok.convert_tokens_to_ids("<|im_end|>")
+        self.reset()
+
+    # ---- state -------------------------------------------------------------------------------------------------
+
+    def reset(self):
+        self.cache = None
+        self.mask = None
+        self.key = None  # (mode, system prompt): a change starts a fresh context
+        self.needs_close = False  # last generation stopped without <|im_end|>
+        self.turn_snapshot = None
+        self._end_audio()
+
+    def _end_audio(self):
+        self.resampler = None
+        self.feature_cache = None
+        self.encoder_cache = None
+        self.audio_in_s = 0.0
+        self.audio_tokens = 0
+
+    @property
+    def length(self) -> int:
+        return 0 if self.cache is None else self.cache.get_seq_length()
+
+    def snapshot(self):
+        return copy.deepcopy((self.cache, self.mask, self.needs_close))
+
+    def restore(self, snap):
+        self.cache, self.mask, self.needs_close = snap
+        self._end_audio()
+
+    # ---- appending -----------------------------------------------------------------------------------------------
+
+    def _append_embeds(self, embeds: torch.Tensor):
+        n = embeds.shape[1]
+        if self.cache is None:
+            from transformers import DynamicCache
+
+            self.cache = DynamicCache(config=self.model.config.text_config)
+            self.mask = torch.ones((1, 0), dtype=torch.long, device=self.model.device)
+        self.mask = torch.cat((self.mask, self.mask.new_ones((1, n))), dim=-1)
+        position_ids = self.mask.cumsum(-1).sub(1)[:, -n:]
+        self.model.model.language_model(
+            inputs_embeds=embeds,
+            past_key_values=self.cache,
+            attention_mask=self.mask,
+            position_ids=position_ids,
+            use_cache=True,
+        )
+
+    def append_text(self, text: str):
+        ids = self.tok(text, add_special_tokens=False, return_tensors="pt").input_ids.to(self.model.device)
+        self._append_embeds(self.model.get_input_embeddings()(ids))
+
+    def close_prefix(self) -> str:
+        """Text that closes the previous assistant turn, if generation left it open."""
+        return "<|im_end|>\n" if self.needs_close else "\n" if self.cache is not None else ""
+
+    # ---- a turn --------------------------------------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def begin_turn(self, key, system_block: str, keep_history: bool) -> dict:
+        """Open a user turn. `system_block` is the rendered `<bos><|im_start|>system…<|im_end|>\\n` prefix."""
+        info = {}
+        if not keep_history or key != self.key or self.length > MAX_CONTEXT:
+            if self.cache is not None:
+                info["context_reset"] = "history off" if not keep_history else "settings changed" if key != self.key else "context full"
+            self.reset()
+            self.key = key
+        self._end_audio()
+        t = time.perf_counter()
+        self.turn_snapshot = self.snapshot()
+        info["snapshot_ms"] = round((time.perf_counter() - t) * 1000, 1)
+        prefix = self.close_prefix() if self.cache is not None else system_block
+        self.append_text(prefix + "<|im_start|>user\n")
+        self.needs_close = False
+        info["context_tokens"] = self.length
+        return info
+
+    @torch.inference_mode()
+    def feed_audio(self, pcm: np.ndarray, rate: int, final: bool = False):
+        """Resample, extract mel frames and prefill whatever audio tokens they complete."""
+        if self.resampler is None:
+            self.resampler = soxr.ResampleStream(rate, self.fe.sampling_rate, 1, dtype="float32")
+        pcm16 = self.resampler.resample_chunk(np.asarray(pcm, dtype=np.float32), last=final)
+        self.audio_in_s += len(pcm16) / self.fe.sampling_rate
+        features, self.feature_cache = self.fe.stream(pcm16, self.feature_cache, is_final=final)
+        mel = features["input_features"].to(self.model.device)
+        if not mel.shape[-1] and not final:
+            return
+        out = self.model.prefill_audio(
+            mel, past_key_values=self.cache, encoder_cache=self.encoder_cache, attention_mask=self.mask, is_final=final
+        )
+        self.audio_tokens += out.attention_mask.shape[1] - self.mask.shape[1]
+        self.cache, self.mask, self.encoder_cache = out.past_key_values, out.attention_mask, out.encoder_cache
+
+    @torch.inference_mode()
+    def generate(self, suffix: str, **kwargs):
+        """Append `suffix` and generate from the cache; the cache then covers the whole reply."""
+        ids = self.tok(suffix, add_special_tokens=False, return_tensors="pt").input_ids.to(self.model.device)
+        mask = torch.cat((self.mask, self.mask.new_ones((1, ids.shape[1]))), dim=-1)
+        try:
+            out = self.model.generate(
+                input_ids=ids, attention_mask=mask, past_key_values=self.cache, return_audio_eos=True,
+                return_dict_in_generate=True, **kwargs,
+            )
+        finally:
+            # the cache was advanced in place, whether generation finished or was interrupted
+            self.mask = torch.ones((1, self.length), dtype=torch.long, device=self.model.device)
+        self.needs_close = int(out.sequences[0, -1]) != self.im_end
+        return out