140 lines
6.1 KiB
Python
140 lines
6.1 KiB
Python
"""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
|