utils.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2026 The OpenBMB Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from dataclasses import dataclass
from typing import Any
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Tuple
from typing import Union

import torch
import torch.nn.functional as F
import torch.nn.utils.parametrize as P
from transformers.cache_utils import DynamicCache

logger = logging.getLogger(__name__)


# text
@dataclass
class GenerateChunkOutput:
    chunk_token_ids: torch.Tensor
    current_inputs_embeds: torch.Tensor
    input_last_hidden_states: Optional[torch.Tensor]  # for tts use_speaker_embedding
    last_hidden_states: Optional[torch.Tensor]  # for tts input feature (projector_semantic)
    past_key_values: Optional[torch.Tensor]
    finished: bool


class ChunkPrefillChunkGenerate:
    def __init__(self, model, tokenizer, terminators):
        self.tokenizer = tokenizer
        self.model = model
        self.terminators = terminators
        self.terminators_ids = [tokenizer.convert_tokens_to_ids(i) for i in self.terminators]
        self.embedding_layer = self.model.get_input_embeddings()

        self.forbidden_tokens = [
            ":",
            "：",
            "；",
            "#",
            "“",
            "”",
            "‘",
            "’",
            "@",
            "*",
            "【",
            "】",
            "「",
            "」",
            "(",
            ")",
            "（",
            "）",
            "[",
            "]",
            "&",
            "/",
            "$",
        ]

        self.forbidden_token_ids = [tokenizer.convert_tokens_to_ids(i) for i in self.forbidden_tokens]
        bad_token_ids = getattr(tokenizer, "bad_token_ids", [])
        if bad_token_ids:
            self.forbidden_token_ids.extend(bad_token_ids)

    @staticmethod
    def prepare_generation_config(do_sample, max_new_tokens=50, min_new_tokens=0, **kwargs):
        num_beams = kwargs.get("num_beams", 3)
        generation_config = {
            "num_beams": num_beams,
            "top_p": 0.8,
            "top_k": 100,
            "temperature": 0.7,
            "do_sample": True,
            "repetition_penalty": 1.05,
        }

        if do_sample:
            generation_config.update(
                {
                    "top_p": 0.8,
                    "top_k": 100,
                    "temperature": 0.7,
                    "do_sample": True,
                    "repetition_penalty": 1.05,
                }
            )
        elif num_beams > 1:
            generation_config.update({"num_beams": num_beams, "repetition_penalty": 1.2, "do_sample": False})
        else:
            generation_config.update({"do_sample": False, "repetition_penalty": 1.05})

        generation_config.update((k, kwargs[k]) for k in generation_config.keys() & kwargs.keys())
        generation_config["min_new_tokens"] = min_new_tokens
        generation_config["max_new_tokens"] = max_new_tokens

        return generation_config

    def chunk_generate(
        self,
        inputs_embeds: torch.Tensor,
        past_key_values,
        is_first_generate_chunk: bool,
        chunk_size: int,
        return_hidden_states: bool,
        do_sample: bool,
        temperature: float,
        top_p: float,
        top_k: int,
        repetition_penalty: float = 1.05,
        length_penalty: float = 1.0,
        all_input_ids: Optional[torch.Tensor] = None,
    ) -> GenerateChunkOutput:
        """
        Args:
            inputs_embeds: [1, seq_len, hidden_dim], Input embeddings of current chunk.
            past_key_values: [num_layers, 2, batch_size, num_heads, seq_len, head_dim], Past key values for llm.
            is_first_generate_chunk: bool, Whether this is the first generate chunk.
            chunk_size: int, The size of the current chunk, default is 10, and it is fixed during training.
            return_hidden_states: bool Whether to return the hidden states, default is True.
            do_sample: bool Whether to sample from the model, default is True.
            temperature: float The temperature for the model, default is 0.7.
            top_p: float The top-p for the model, default is 0.8.
            top_k: int The top-k for the model, default is 100.
            repetition_penalty: float, The repetition penalty for the model, default is 1.05.
            length_penalty: float, The length penalty for the model, default is 1.0. Higher value means more detailed generation.
            all_input_ids: Optional[torch.Tensor], The input ids for the current chunk.
        """

        finished = False
        current_inputs_embeds = inputs_embeds.clone()
        input_last_hidden_states = []
        last_hidden_states = []
        generated_tokens = []

        for token_idx in range(chunk_size):
            if is_first_generate_chunk and token_idx == 0:
                # first generate chunk, prefill inputs_embeds
                model_inputs = {
                    "inputs_embeds": current_inputs_embeds,
                    "past_key_values": past_key_values,
                    "use_cache": True,
                    "output_hidden_states": return_hidden_states,
                }
            else:  # for all other cases: prefill the latest generated token
                model_inputs = {
                    "inputs_embeds": current_inputs_embeds[:, -1:, :],
                    "past_key_values": past_key_values,
                    "use_cache": True,
                    "output_hidden_states": return_hidden_states,
                }

            with torch.no_grad():
                outputs = self.model(**model_inputs)

            # last token's logits
            logits = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=inputs_embeds.device)

            # forbid specific tokens decoding = model.generate@suppress_tokens
            if self.forbidden_token_ids:
                logits[:, self.forbidden_token_ids] = float("-inf")

            past_key_values = outputs.past_key_values

            PENALTY_WINDOW_SIZE = 128

            # apply repetition penalty
            if repetition_penalty != 1.0:
                # get token ids for repetition penalty
                if all_input_ids is not None:
                    # use global input ids (including original input and generated part)
                    if len(generated_tokens) > 0:
                        generated_token_ids = torch.cat(generated_tokens, dim=1)
                        current_sequence = torch.cat(
                            [
                                all_input_ids[:, -PENALTY_WINDOW_SIZE:],
                                generated_token_ids,
                            ],
                            dim=1,
                        )
                    else:
                        current_sequence = all_input_ids[:, -PENALTY_WINDOW_SIZE:]
                    unique_token_ids = torch.unique(current_sequence.squeeze(0))
                elif len(generated_tokens) > 0:
                    # revert to original logic: only use generated tokens
                    generated_token_ids = torch.cat(generated_tokens, dim=1).squeeze(0)
                    unique_token_ids = torch.unique(generated_token_ids)
                else:
                    unique_token_ids = torch.tensor([], dtype=torch.long, device=logits.device)

                # apply repetition penalty
                for token_id in unique_token_ids:
                    if logits[0, token_id] > 0:
                        logits[0, token_id] = logits[0, token_id] / repetition_penalty
                    else:
                        logits[0, token_id] = logits[0, token_id] * repetition_penalty

            # apply length penalty, higher value means more detailed generation
            if length_penalty != 1.0:
                for eos_token_id in self.terminators_ids:
                    if logits[0, eos_token_id] > 0:
                        logits[0, eos_token_id] = logits[0, eos_token_id] / length_penalty
                    else:
                        logits[0, eos_token_id] = logits[0, eos_token_id] * length_penalty

            # apply temperature
            if temperature != 1.0:
                logits = logits / temperature

            if do_sample:
                # Top-k filtering
                if top_k > 0:
                    top_k_logits, top_k_indices = torch.topk(logits, min(top_k, logits.size(-1)))
                    logits_filtered = torch.full_like(logits, float("-inf"))
                    logits_filtered.scatter_(1, top_k_indices, top_k_logits)
                    logits = logits_filtered

                # Top-p filtering
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

                    # remove tokens with cumulative probability greater than top_p
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                    sorted_indices_to_remove[..., 0] = 0

                    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
                    logits[indices_to_remove] = float("-inf")

                # sampling
                probs = F.softmax(logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(logits, dim=-1, keepdim=True)

            if return_hidden_states:
                if is_first_generate_chunk and token_idx == 0:
                    input_last_hidden_states.append(outputs.hidden_states[-1])
                else:
                    last_hidden_states.append(outputs.hidden_states[-1])

            # if terminator token, stop generating
            if next_token.item() in self.terminators_ids:
                finished = True
                break

            generated_tokens.append(next_token)

            # convert new token to embeddings and concatenate
            next_token_embed = self.embedding_layer(next_token)

            # update inputs_embeds, add one
            current_inputs_embeds = torch.cat([current_inputs_embeds, next_token_embed], dim=1)

        if len(generated_tokens) > 0:
            chunk_token_ids = torch.cat(generated_tokens, dim=1)
        else:
            # special case: if last chunk and first predict is eos token, return last token of previous chunk. return a tensor with shape (1, 0)
            if finished:
                chunk_token_ids = torch.zeros((1, 0), dtype=torch.long, device=current_inputs_embeds.device)
            else:
                raise Exception("this should not happen")

        if len(last_hidden_states) > 0:
            last_hidden_states = torch.cat(last_hidden_states, dim=1)
        else:
            # special case: if last chunk, return last token of previous chunk.
            if finished:
                last_hidden_states = torch.cat(last_hidden_states, dim=1)
            else:
                raise Exception("this should not happen")

        if len(input_last_hidden_states) > 0:
            input_last_hidden_states = torch.cat(input_last_hidden_states, dim=1)
        else:
            input_last_hidden_states = None

        return GenerateChunkOutput(
            chunk_token_ids=chunk_token_ids,
            current_inputs_embeds=current_inputs_embeds,
            input_last_hidden_states=input_last_hidden_states,
            last_hidden_states=last_hidden_states,
            past_key_values=past_key_values,
            finished=finished,
        )


def streaming_token_decoder(token_iterator, tokenizer, skip_special_tokens=False):
    """
    Incrementally decode tokens from an iterator, handling partial multi-byte characters.

    When streaming tokens, multi-byte characters (like Chinese) may be split across multiple
    tokens. Decoding partial tokens results in replacement characters (U+FFFD). This function
    buffers tokens and only yields complete characters.

    Args:
        token_iterator: An iterator yielding (token_ids, is_finished) tuples.
                       token_ids can be torch.Tensor or any iterable of integers.
        tokenizer: The tokenizer to use for decoding.
        skip_special_tokens: Whether to skip special tokens during decoding.

    Yields:
        (decoded_text, is_finished) tuples where decoded_text is the new text since last yield.
    """
    accumulated_token_ids = []
    yielded_text_len = 0

    for token_ids, is_finished in token_iterator:
        # Accumulate token IDs
        if torch.is_tensor(token_ids):
            accumulated_token_ids.extend(token_ids.reshape(-1).tolist())
        else:
            accumulated_token_ids.extend(list(token_ids) if hasattr(token_ids, "__iter__") else [token_ids])

        # Decode all accumulated tokens
        full_decoded = tokenizer.decode(accumulated_token_ids, skip_special_tokens=skip_special_tokens)

        if is_finished:
            # Final chunk - yield all remaining text
            new_text = full_decoded[yielded_text_len:]
            yield new_text, is_finished
        else:
            # Find safe prefix without incomplete multi-byte characters
            # The replacement character '�' (U+FFFD) indicates incomplete decoding
            new_text = full_decoded[yielded_text_len:]

            # Hold back text ending with replacement character (incomplete UTF-8 sequence)
            safe_end = len(new_text)
            while safe_end > 0 and new_text[safe_end - 1] == "\ufffd":
                safe_end -= 1

            safe_text = new_text[:safe_end] if safe_end > 0 else ""
            yielded_text_len += len(safe_text)
            yield safe_text, is_finished


def torch_clone_recursive(obj):
    """Recursively clone nested containers of torch.Tensors.

    Supported container types: dict, list, tuple. Non-container non-Tensor
    objects are returned as-is.
    """
    if torch.is_tensor(obj):
        return obj.clone()
    elif isinstance(obj, dict):
        return {k: torch_clone_recursive(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [torch_clone_recursive(v) for v in obj]
    elif isinstance(obj, tuple):
        return tuple(torch_clone_recursive(v) for v in obj)
    else:
        raise ValueError(f"Unsupported type: {type(obj)}")


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Rotate half the hidden dims of the input for RoPE."""
    dim = x.shape[-1]
    x1 = x[..., : dim // 2]
    x2 = x[..., dim // 2 :]
    return torch.cat((-x2, x1), dim=-1)


@dataclass
class SpeculativeSnapshot:
    """Speculative snapshot for VAD speculative rollback.

    Used in VAD speculative execution: creates a snapshot after streaming_prefill
    and before streaming_generate. If speculation fails (user continues speaking),
    the state can be restored to continue streaming_prefill.

    Implementation:
    - LLM KV Cache: only record length, restore by truncation (zero extra VRAM)
    - Audio KV Cache: requires cloning, as generate sets it to None
    - Mel processor: save full state snapshot (including buffer)
    """

    # KV Cache length (for truncation recovery)
    llm_cache_length: int
    audio_cache_length: int

    # session state
    new_user_msg: bool
    llm_generated: bool
    llm_generate_completed: bool

    # Round management
    next_round_id: int
    pending_round_id: Optional[int]
    omni_chunk_history_length: int

    # TTS state (requires cloning, but usually small)
    tts_last_turn_tokens: Optional[torch.Tensor]

    # Streaming processor state
    audio_chunk_idx: int

    # Mel processor state snapshot (including buffer)
    mel_processor_snapshot: Optional[dict] = None

    # Audio encoder KV cache (requires cloning to ensure determinism after recovery)
    audio_past_key_values: Optional[tuple] = None

    # timestamp (for debugging)
    timestamp: float = 0.0

    # debug field: for verifying correctness of recovery
    llm_cache_checksum: Optional[float] = None  # LLM KV Cache first layer K sum
    audio_cache_checksum: Optional[float] = None  # Audio KV Cache first layer K sum
    mel_buffer_checksum: Optional[float] = None  # Mel buffer sum

    # RNG state (key: for ensuring determinism of dithering etc. after recovery)
    rng_state_cpu: Optional[torch.Tensor] = None  # torch CPU RNG state
    rng_state_cuda: Optional[torch.Tensor] = None  # torch CUDA RNG state (if on GPU)

    def summary(self) -> str:
        mel_buf_len = 0
        if self.mel_processor_snapshot:
            buf = self.mel_processor_snapshot.get("buffer")
            if buf is not None:
                mel_buf_len = len(buf)
        return (
            f"llm_cache={self.llm_cache_length}, "
            f"audio_cache={self.audio_cache_length}, "
            f"audio_chunk_idx={self.audio_chunk_idx}, "
            f"mel_buffer={mel_buf_len}, "
            f"history_len={self.omni_chunk_history_length}, "
            f"new_user_msg={self.new_user_msg}, "
            f"llm_generated={self.llm_generated}"
        )


# tts
@dataclass
class TTSSamplingParams:
    top_p: float = 0.85
    min_p: float = 0.01
    top_k: int = 25
    repetition_penalty: float = 1.05
    temperature: float = 0.8
    win_size: int = 16
    tau_r: float = 0.1


class TTSStreamingGenerator:
    """
    Streaming generator for TTS that processes chunks and yields audio tokens in real-time.

    Supported attention types:
    - full_attention: Full attention, all tokens can attend to each other
    - sliding_window: Sliding window attention, KV cache is truncated to fixed size (token_window_size)
    - sliding_recompute: Sliding recompute, only keep previous chunk and recompute with current chunk
    - reindex: Keep first chunk as sink, reindex sliding window positions via RoPE rotation
    """

    def __init__(
        self,
        model,
        temperature: float,
        eos_token: Union[int, torch.Tensor],
        chunk_size: int = 25,  # s3tokenizer 1s = 25token
        tts_last_turn_tokens: torch.Tensor = None,
        logits_processors=None,
        logits_warpers=None,
    ):
        self.tts = model
        self.device = model.device
        self.temperature = torch.tensor([temperature], dtype=torch.float, device=self.device)
        self.eos_token = (
            torch.tensor(eos_token, device=self.device) if isinstance(eos_token, int) else eos_token.to(self.device)
        )

        self.num_vq = model.num_vq
        self.num_audio_tokens = model.num_audio_tokens
        self.recomputed_chunks = model.recomputed_chunks
        self.emb_code = model.emb_code
        self.head_code = model.head_code

        # Attention type and window sizes
        self.attention_type = model.attention_type  # "full_attention", "sliding_window", "sliding_recompute", "reindex"
        self.chunk_window_size = model.chunk_window_size  # chunk-level window for sliding_recompute (default 2)
        self.token_window_size = model.token_window_size  # token-level window for sliding_window/reindex (default 300)

        # RoPE config (for reindex mode)
        self.rope_theta = model.model.config.rope_theta
        self.head_dim = model.model.config.hidden_size // model.model.config.num_attention_heads

        # Logits processors
        self.logits_processors = logits_processors if logits_processors is not None else []
        # Logits warpers (like TopP/TopK), separate from processors
        self.logits_warpers = logits_warpers if logits_warpers is not None else []

        # initialize state
        self.past_key_values = None
        self.text_start_pos = 0
        self.idx = -1  # start from -1, become 0 when first called
        self.all_conditions = []
        self.all_generated_tokens = []
        self.tts_last_turn_tokens = tts_last_turn_tokens
        self.spk_emb = None

        audio_bos = [self.tts.audio_bos_token_id]
        audio_bos = torch.Tensor(audio_bos).to(self.tts.emb_text.weight.device, dtype=torch.long)

        self.audio_bos_embeds = self.tts.emb_text(audio_bos).unsqueeze(0)
        self.text_eos_embed = self.tts.emb_text(
            torch.tensor(
                [self.tts.config.text_eos_token_id],
                device=self.tts.emb_text.weight.device,
                dtype=torch.long,
            )
        ).unsqueeze(0)

        # buffer related, used to fill up chunk_size and yield to outside
        self.chunk_size = chunk_size
        self._token_buffer: List[torch.Tensor] = []

        # Chunk info tracking for sliding_recompute and reindex
        self._chunk_info: List[dict] = []
        self._total_seq_len = 0

        # Reindex mode: track sink (first chunk) length
        self._sink_kv_len = 0

    def _build_recompute_inputs(self, current_condition: torch.Tensor) -> torch.Tensor:
        """Build recompute inputs for sliding_recompute mode."""
        if len(self._chunk_info) == 0:
            return current_condition

        prev_chunk = self._chunk_info[-1]
        prev_condition = prev_chunk["condition"]
        prev_audio_tokens = prev_chunk["audio_tokens"]

        recompute_list = [prev_condition]
        if len(prev_audio_tokens) > 0:
            prev_audio_embeds = torch.cat([self.emb_code[0](tok) for tok in prev_audio_tokens], dim=1)
            recompute_list.append(prev_audio_embeds)

        recompute_list.append(current_condition)
        return torch.cat(recompute_list, dim=1)

    def _truncate_kv_cache_sliding_window(self):
        """Truncate KV cache for sliding_window mode."""
        if self.past_key_values is None:
            return

        if hasattr(self.past_key_values, "get_seq_length"):
            current_kv_len = self.past_key_values.get_seq_length()
        else:
            current_kv_len = self.past_key_values[0][0].shape[2]

        if current_kv_len <= self.token_window_size:
            return

        new_cache = DynamicCache()
        num_layers = (
            len(self.past_key_values.key_cache)
            if hasattr(self.past_key_values, "key_cache")
            else len(self.past_key_values)
        )

        for layer_idx in range(num_layers):
            if hasattr(self.past_key_values, "key_cache"):
                key = self.past_key_values.key_cache[layer_idx][:, :, -self.token_window_size :, :]
                value = self.past_key_values.value_cache[layer_idx][:, :, -self.token_window_size :, :]
            else:
                key = self.past_key_values[layer_idx][0][:, :, -self.token_window_size :, :]
                value = self.past_key_values[layer_idx][1][:, :, -self.token_window_size :, :]
            new_cache.update(key, value, layer_idx)

        self.past_key_values = new_cache

    @staticmethod
    def _apply_rope_rotation(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        """Apply RoPE rotation to tensor."""
        return x * cos + rotate_half(x) * sin

    def _compute_rope_cos_sin(self, positions: torch.Tensor, device: torch.device, dtype: torch.dtype):
        """Compute RoPE cos and sin for given positions."""
        dim_half = self.head_dim // 2
        freq_seq = torch.arange(0, dim_half, dtype=torch.float32, device=device)
        inv_freq = 1.0 / (self.rope_theta ** (freq_seq / dim_half))

        # positions: [seq_len]
        angles = positions.float().unsqueeze(-1) * inv_freq.unsqueeze(0)  # [seq_len, dim_half]
        angles = torch.cat([angles, angles], dim=-1)  # [seq_len, head_dim]

        cos = angles.cos().to(dtype)
        sin = angles.sin().to(dtype)
        return cos, sin

    def _reindex_kv_cache(self):
        """
        Reindex KV cache for reindex mode:
        1. Keep first chunk as attention sink
        2. Keep last chunk
        3. Discard middle chunks
        4. Reindex the last chunk's key positions to be right after sink via RoPE rotation
        """
        if self.past_key_values is None or len(self._chunk_info) < 2:
            return

        # Get current KV cache length
        if hasattr(self.past_key_values, "get_seq_length"):
            current_kv_len = self.past_key_values.get_seq_length()
        else:
            current_kv_len = self.past_key_values[0][0].shape[2]

        # Calculate sink length (first chunk)
        sink_len = self._chunk_info[0]["condition_len"] + self._chunk_info[0]["audio_token_count"]

        # Last chunk length
        last_chunk = self._chunk_info[-1]
        last_chunk_len = last_chunk["condition_len"] + last_chunk["audio_token_count"]

        keep_len = sink_len + last_chunk_len

        # Get device and dtype
        device = self.past_key_values.key_cache[0].device
        dtype = self.past_key_values.key_cache[0].dtype

        if current_kv_len <= keep_len:
            last_chunk_kv_len = current_kv_len - sink_len
            if last_chunk_kv_len <= 0:
                return
            self.text_start_pos = current_kv_len
            return

        # Step 1: Truncate KV cache - keep sink and last chunk
        new_cache = DynamicCache()
        num_layers = len(self.past_key_values.key_cache)

        original_start_pos = current_kv_len - last_chunk_len
        new_start_pos = sink_len
        delta = new_start_pos - original_start_pos  # This is a scalar constant
        delta_positions = torch.full((last_chunk_len,), delta, dtype=torch.float32, device=device)

        # Compute rotation cos/sin
        cos, sin = self._compute_rope_cos_sin(delta_positions, device, dtype)
        cos = cos.unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, head_dim]
        sin = sin.unsqueeze(0).unsqueeze(0)

        for layer_idx in range(num_layers):
            key_full = self.past_key_values.key_cache[layer_idx]
            value_full = self.past_key_values.value_cache[layer_idx]

            # Extract sink and last chunk
            key_sink = key_full[:, :, :sink_len, :]
            value_sink = value_full[:, :, :sink_len, :]
            key_last = key_full[:, :, -last_chunk_len:, :]
            value_last = value_full[:, :, -last_chunk_len:, :]

            # Apply RoPE rotation to reindex key positions
            key_last_reindexed = self._apply_rope_rotation(key_last, cos, sin)

            # Concatenate sink and reindexed last chunk
            key = torch.cat([key_sink, key_last_reindexed], dim=2)
            value = torch.cat([value_sink, value_last], dim=2)

            new_cache.update(key, value, layer_idx)

        self.past_key_values = new_cache

        # Update text_start_pos to reflect new positions
        self.text_start_pos = sink_len + last_chunk_len

    @torch.inference_mode()
    def generate_with_buffer(
        self,
        condition: torch.Tensor,
        text_finished: bool = False,
        max_new_token: int = 500,
    ):
        """input a condition embedding chunk, generate audio token each time,
        and accumulate to buffer, only yield when buffer satisfies chunk_size.

        Yields:
            torch.Tensor of shape [chunk_size] (2D: [1, chunk_size])
        """
        self.idx += 1
        self.device = self.tts.device

        # if text finished, first concatenate Text EOS
        if text_finished:
            condition = torch.cat([condition, self.text_eos_embed], dim=1)

        # always concatenate Audio BOS
        condition = torch.cat([condition, self.audio_bos_embeds], dim=1).to(self.device)

        self.all_conditions.append(condition)

        # Initialize current chunk info
        current_chunk_info = {
            "condition_len": condition.shape[1],
            "audio_token_count": 0,
            "condition": condition.clone(),
            "audio_tokens": [],
        }

        # Handle different attention types
        if self.attention_type == "sliding_recompute" and self.idx >= 1:
            # sliding_recompute: discard KV cache, recompute with previous + current chunk
            self.past_key_values = None
            current_condition = self._build_recompute_inputs(condition)
            self.text_start_pos = 0
        elif self.attention_type == "reindex" and self.idx >= 1:
            # reindex: truncate KV cache keeping sink + last chunk, reindex positions via RoPE
            self._reindex_kv_cache()
            current_condition = condition
            # Always update text_start_pos based on actual KV cache length (like reference code)
            if self.past_key_values is not None:
                if hasattr(self.past_key_values, "get_seq_length"):
                    kv_len = self.past_key_values.get_seq_length()
                else:
                    kv_len = self.past_key_values[0][0].shape[2]
                self.text_start_pos = kv_len
        else:
            current_condition = condition

        condition_length = current_condition.shape[1]
        prefill_len = condition_length
        finished = torch.zeros(1, dtype=torch.bool, device=self.device)
        chunk_generated_tokens = []

        for t in range(max_new_token):
            if t == 0:
                inputs_embeds = current_condition
                pos_ids = torch.arange(
                    self.text_start_pos,
                    self.text_start_pos + condition_length,
                    dtype=torch.long,
                    device=self.device,
                ).unsqueeze(0)
            else:
                last = self.all_generated_tokens[-1]
                # last: [1,1], directly as code id
                inputs_embeds = self.emb_code[0](last)
                pos_ids = torch.tensor(
                    [self.text_start_pos + prefill_len + t - 1],
                    dtype=torch.long,
                    device=self.device,
                ).unsqueeze(0)

            outputs = self.tts.model(
                position_ids=pos_ids,
                past_key_values=self.past_key_values,
                inputs_embeds=inputs_embeds,
                use_cache=True,
            )
            hidden_states = outputs.last_hidden_state

            # Handle KV cache based on attention type
            if self.attention_type == "sliding_window":
                self.past_key_values = outputs.past_key_values
                self._truncate_kv_cache_sliding_window()
            else:
                self.past_key_values = outputs.past_key_values

            with P.cached():
                logits = torch.empty(
                    hidden_states.size(0),
                    hidden_states.size(1),
                    self.num_audio_tokens,
                    self.num_vq,
                    dtype=torch.float,
                    device=self.device,
                )
                for num_vq_iter in range(self.num_vq):
                    x: torch.Tensor = self.head_code[num_vq_iter](hidden_states)
                    logits[..., num_vq_iter] = x
                    del x

            del hidden_states

            logits = logits[:, -1].float()

            logits = logits.permute(0, 2, 1)
            logits = logits.reshape(-1, logits.size(2))

            logits /= self.temperature

            audio_bos = len(self.all_generated_tokens) == 0 and t == 0

            if not audio_bos:
                # use generated tokens (current chunk) as input for processor/warper (align with modeling_minicpmo)
                all_generated_tokens = torch.cat(self.all_generated_tokens, dim=1).to(self.device)  # [1, T]
                for processor in self.logits_processors:
                    logits = processor(all_generated_tokens, logits)

                for warper in self.logits_warpers:
                    logits = warper(all_generated_tokens, logits)
                del all_generated_tokens

            # sample next token (only use first codebook, same as generate)
            scores = F.softmax(logits, dim=-1)
            idx_next = torch.multinomial(scores, num_samples=1)  # [(B*num_vq), 1]
            next_id = idx_next.view(-1, self.num_vq)[:, 0:1]  # only take first codebook → [B, 1]
            del scores

            if next_id.eq(
                self.eos_token
            ).any():  # generated audio eos token, means this chunk is finished, no longer generate new tokens
                finished[:] = True
            else:  # eos token cannot be added to buffer, he does not speak.
                # convert next_id to correct shape [1, 1], no num_vq dimension
                if next_id.dim() == 0:  # if scalar
                    next_tok = next_id.unsqueeze(0).unsqueeze(0)  # [1, 1]
                elif next_id.dim() == 1:  # if 1D [1]
                    next_tok = next_id.unsqueeze(0)  # [1, 1]
                else:
                    next_tok = next_id

                self.all_generated_tokens.append(next_tok)
                chunk_generated_tokens.append(next_tok)

                # Update chunk info for sliding_recompute
                current_chunk_info["audio_tokens"].append(next_tok.clone())
                current_chunk_info["audio_token_count"] += 1

                self._token_buffer.append(next_tok)

            if len(self._token_buffer) == 0:
                # case 1: if last text chunk, yield None
                if text_finished:
                    yield torch.empty(1, 0, dtype=torch.long, device=self.device), True
                    break
                # case 2: if not last text chunk, break directly
                else:
                    break
            else:  # buffer has something
                # case 1: if buffer is larger/equal to chunk_size, yield out
                if len(self._token_buffer) >= self.chunk_size:
                    batch = torch.cat(self._token_buffer[: self.chunk_size], dim=1)  # [1, chunk_size]
                    yield batch, False  # → [1, chunk_size]
                    # discard yielded part
                    self._token_buffer = self._token_buffer[self.chunk_size :]

                # case 2: if buffer is smaller than chunk_size
                else:
                    # if generation finished, and is the last text chunk, yield all remaining tokens, then break
                    if finished.all():
                        if text_finished:
                            batch = torch.cat(self._token_buffer, dim=1)  # [1, chunk_size]
                            yield batch, True  # → [1, chunk_size]
                            self._token_buffer = []
                            break
                        else:
                            # not the last text chunk, need to wait for next text chunk to fill up buffer, then this call ends
                            break
                    else:  # generation of this audio chunk is not finished, continue generating
                        continue

        # Save current chunk info for sliding_recompute and reindex
        self._chunk_info.append(current_chunk_info)
        self._total_seq_len += condition.shape[1] + len(chunk_generated_tokens)

        # Update text_start_pos based on attention type
        if self.attention_type == "sliding_recompute":
            # sliding_recompute: will be reset at next chunk start, update normally here
            self.text_start_pos += prefill_len + len(chunk_generated_tokens)
        elif self.attention_type == "reindex":
            # reindex: position based on actual KV cache length (positions have been reindexed to be continuous)
            if self.past_key_values is not None:
                if hasattr(self.past_key_values, "get_seq_length"):
                    self.text_start_pos = self.past_key_values.get_seq_length()
                else:
                    self.text_start_pos = self.past_key_values[0][0].shape[2]
            else:
                self.text_start_pos += condition.shape[1] + len(chunk_generated_tokens)
        else:
            self.text_start_pos += condition.shape[1] + len(chunk_generated_tokens)
        # note: remaining tokens in buffer will be kept, and accumulated next time


# sliding window
@dataclass
class StreamingWindowConfig:
    text_window_high_tokens: int = 8000
    text_window_low_tokens: int = 6000


@dataclass
class DuplexWindowConfig:
    """duplex sliding window configuration

    sliding window mode:
    - "off": disable sliding window
    - "basic": basic sliding window (trigger by cache length)
    - "context": sliding window with context (trigger by unit number, preserve generated text to previous)
    """

    # sliding window mode
    sliding_window_mode: str = "off"  # "off" / "basic" / "context"

    # basic sliding window parameters
    basic_window_high_tokens: int = 8000  # high watermark: trigger sliding window when exceeded
    basic_window_low_tokens: int = 6000  # low watermark: keep to this value after sliding window

    # context sliding window parameters
    context_previous_max_tokens: int = 500  # previous maximum token number
    context_max_units: int = 24  # maximum unit number (trigger sliding window when exceeded)

    # verification mode (for comparison test)
    verify_mode: bool = False  # whether to enable verification log


def as_dynamic_cache(past_key_values):
    """Convert legacy tuple cache to DynamicCache if needed."""
    if isinstance(past_key_values, DynamicCache):
        return past_key_values

    if isinstance(past_key_values, tuple):
        return DynamicCache.from_legacy_cache(past_key_values)

    return past_key_values


def get_kv_cache_length(cache) -> int:
    """Get the sequence length of a KV cache.

    Args:
        cache: DynamicCache or tuple-based cache

    Returns:
        The number of tokens in the cache
    """
    if cache is None:
        return 0

    if isinstance(cache, DynamicCache):
        if not cache.key_cache or not cache.key_cache[0].numel():
            return 0
        return cache.key_cache[0].shape[-2]

    if isinstance(cache, tuple):
        return cache[0][0].shape[2]

    return 0


def get_rotary_cos_sin(
    head_dim: int,
    positions: torch.Tensor,
    device: torch.device,
    dtype: torch.dtype,
    rope_theta: float = 10000.0,
    inv_freq_cache: Optional[Dict[Tuple, torch.Tensor]] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Compute RoPE cos and sin components for given positions.

    Args:
        head_dim: Dimension of each attention head
        positions: Position indices tensor
        device: Target device
        dtype: Target dtype
        rope_theta: RoPE base frequency (default 10000.0)
        inv_freq_cache: Optional cache dict for inverse frequencies

    Returns:
        Tuple of (cos, sin) tensors with shape [1, 1, seq_len, head_dim]
    """
    cache_key = (head_dim, device)

    inv_freq = inv_freq_cache.get(cache_key) if inv_freq_cache is not None else None
    if inv_freq is None or inv_freq.device != device or inv_freq.shape[0] != head_dim // 2:
        exponent = torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim
        inv_freq = 1.0 / (rope_theta**exponent)
        if inv_freq_cache is not None:
            inv_freq_cache[cache_key] = inv_freq

    positions = positions.to(device=device, dtype=torch.float32)
    angles = torch.einsum("i,j->ij", positions, inv_freq)
    cos = torch.cos(angles)
    sin = torch.sin(angles)

    # Use cat instead of repeat_interleave, consistent with model's original RotaryEmbedding
    # Original: emb = torch.cat((freqs, freqs), dim=-1) -> [f0, f1, ..., f_{d/2}, f0, f1, ..., f_{d/2}]
    cos_full = torch.cat([cos, cos], dim=-1).to(dtype=dtype)
    sin_full = torch.cat([sin, sin], dim=-1).to(dtype=dtype)
    cos_full = cos_full.unsqueeze(0).unsqueeze(0)
    sin_full = sin_full.unsqueeze(0).unsqueeze(0)
    return cos_full, sin_full


def realign_rotary_suffix(
    suffix_keys: torch.Tensor,
    old_positions: torch.Tensor,
    new_positions: torch.Tensor,
    rope_theta: float = 10000.0,
    inv_freq_cache: Optional[Dict[Tuple, torch.Tensor]] = None,
) -> torch.Tensor:
    """Realign RoPE position encoding after cache eviction.

    When tokens are dropped from the middle of a cache, the suffix tokens
    need their RoPE embeddings recalculated with new position indices.

    Args:
        suffix_keys: Key tensor to realign, shape [batch, heads, seq_len, head_dim]
        old_positions: Original position indices
        new_positions: New position indices after eviction
        rope_theta: RoPE base frequency
        inv_freq_cache: Optional cache dict for inverse frequencies

    Returns:
        Realigned key tensor with same shape as input
    """
    if suffix_keys.numel() == 0:
        return suffix_keys

    head_dim = suffix_keys.shape[-1]
    device = suffix_keys.device
    dtype = suffix_keys.dtype

    # Compute old position cos/sin
    cos_old, sin_old = get_rotary_cos_sin(head_dim, old_positions, device, dtype, rope_theta, inv_freq_cache)

    # Inverse transform: recover original key
    base = cos_old * suffix_keys - sin_old * rotate_half(suffix_keys)

    # Compute new position cos/sin
    cos_new, sin_new = get_rotary_cos_sin(head_dim, new_positions, device, dtype, rope_theta, inv_freq_cache)

    # Forward transform: re-encode with new positions
    return cos_new * base + sin_new * rotate_half(base)


def drop_tokens_from_cache(
    cache: Optional[DynamicCache | Tuple],
    length: int,
    preserve: int,
    position_offset: int,
    rope_theta: float = 10000.0,
    inv_freq_cache: Optional[Dict[Tuple, torch.Tensor]] = None,
) -> Tuple[Optional[DynamicCache], int, bool]:
    """Drop tokens from a KV cache while preserving system prompt.

    Removes tokens in the range [preserve, preserve + length) from the cache,
    realigning RoPE embeddings for the suffix.

    Args:
        cache: DynamicCache or tuple-based cache (will be converted to DynamicCache)
        length: Number of tokens to drop
        preserve: Number of tokens to preserve at the start (system prompt)
        position_offset: Current position offset for RoPE calculation
        rope_theta: RoPE base frequency
        inv_freq_cache: Optional cache dict for inverse frequencies

    Returns:
        Tuple of (cache, new_position_offset, success)
        Note: Tuple cache will be converted to DynamicCache. Modification is in-place.
    """
    if cache is None or length <= 0:
        return cache, position_offset, False

    cache = as_dynamic_cache(cache)

    total_len = get_kv_cache_length(cache)
    if total_len <= 0:
        return cache, position_offset, False

    preserve = min(preserve, total_len)
    available = total_len - preserve

    if available < length:
        logger.warning(
            "Cannot drop %d tokens: only %d available (total=%d, preserve=%d)",
            length,
            available,
            total_len,
            preserve,
        )
        return cache, position_offset, False

    suffix_len = total_len - preserve - length
    # note: after RoPE reindex, the position of cache has been compressed (from preserve start)
    # so here should not add position_offset, but use the actual layout of current cache
    suffix_offset = preserve + length  # suffix current position in cache
    prefix_offset = preserve  # suffix new position (follow preserve)

    # Prepare position tensors for RoPE realignment
    old_positions = None
    new_positions = None
    if suffix_len > 0:
        device = cache.key_cache[0].device
        old_positions = torch.arange(
            suffix_offset,
            suffix_offset + suffix_len,
            device=device,
            dtype=torch.long,
        )
        new_positions = torch.arange(
            prefix_offset,
            prefix_offset + suffix_len,
            device=device,
            dtype=torch.long,
        )

    keep_len = total_len - length

    # Process each layer (in-place modification)
    for layer_idx in range(len(cache.key_cache)):
        key_tensor = cache.key_cache[layer_idx]
        value_tensor = cache.value_cache[layer_idx]

        if not key_tensor.numel():
            continue

        # Preserve prefix (system prompt)
        prefix_keys = key_tensor[:, :, :preserve, :]
        prefix_values = value_tensor[:, :, :preserve, :]

        if suffix_len > 0:
            # Keep and realign suffix
            suffix_keys = key_tensor[:, :, preserve + length :, :]
            suffix_values = value_tensor[:, :, preserve + length :, :]

            if old_positions is not None and new_positions is not None and suffix_keys.numel():
                suffix_keys = realign_rotary_suffix(
                    suffix_keys,
                    old_positions,
                    new_positions,
                    rope_theta,
                    inv_freq_cache,
                )

            cache.key_cache[layer_idx] = torch.cat([prefix_keys, suffix_keys], dim=-2).contiguous()
            cache.value_cache[layer_idx] = torch.cat([prefix_values, suffix_values], dim=-2).contiguous()
        else:
            cache.key_cache[layer_idx] = prefix_keys.contiguous()
            cache.value_cache[layer_idx] = prefix_values.contiguous()

    cache.crop(keep_len)
    cache._seen_tokens = max(keep_len, 0)

    new_offset = position_offset + length
    logger.debug("Dropped %d tokens from cache, new length=%d", length, keep_len)

    return cache, new_offset, True


# stream decoder
def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float("inf")):
    logits = logits.clone()

    # Top-k filtering
    if top_k > 0:
        top_k = min(top_k, logits.size(-1))
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    # Top-p (nucleus) filtering
    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        probs = F.softmax(sorted_logits, dim=-1)
        cumulative_probs = torch.cumsum(probs, dim=-1)

        sorted_indices_to_remove = cumulative_probs > top_p
        # keep the first token that exceeds top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[0, indices_to_remove] = filter_value

    return logits


class StreamDecoder:
    def __init__(self, llm, tokenizer, special_token_ids=None, forbidden_token_ids=None):
        self.m = llm
        self.tokenizer = tokenizer
        self.listen_id = self.tokenizer.eos_token_id

        self.chunk_eos_id = self.tokenizer.convert_tokens_to_ids("<|chunk_eos|>")
        self.chunk_tts_eos_id = self.tokenizer.convert_tokens_to_ids("<|chunk_tts_eos|>")
        self.turn_eos_id = self.tokenizer.convert_tokens_to_ids("<|turn_eos|>")
        self.speak_id = self.tokenizer.convert_tokens_to_ids("<|speak|>")

        self.special_token_ids = special_token_ids if special_token_ids is not None else []

        # cache special tokens (used for context sliding window filtering)
        self._all_special_ids = set()
        self._all_special_tokens_text = set()
        if self.tokenizer:
            if hasattr(self.tokenizer, "all_special_ids"):
                self._all_special_ids = set(self.tokenizer.all_special_ids)
            if hasattr(self.tokenizer, "all_special_tokens"):
                self._all_special_tokens_text = set(self.tokenizer.all_special_tokens)

        custom_special_tokens = [
            "<unit>",
            "</unit>",
            "<image>",
            "</image>",
            "<slice>",
            "</slice>",
            "<|listen|>",
            "<|speak|>",
            "<|tts_bos|>",
            "<|tts_eos|>",
            "<|audio_start|>",
            "<|audio_end|>",
            "<|chunk_eos|>",
            "<|chunk_tts_eos|>",
            "<|turn_eos|>",
            "<|audio_start|>",
            "<|audio_end|>",
        ]
        self._all_special_tokens_text.update(custom_special_tokens)
        for token in custom_special_tokens:
            token_id = self.tokenizer.convert_tokens_to_ids(token)
            if token_id is not None and token_id != self.tokenizer.unk_token_id:
                self._all_special_ids.add(token_id)

        if forbidden_token_ids is None:
            self.forbidden_token_ids = []
        elif isinstance(forbidden_token_ids, int):
            self.forbidden_token_ids = [self.forbidden_token_ids]
        else:
            self.forbidden_token_ids = forbidden_token_ids
        self.forbidden_token_ids.append(self.chunk_eos_id)

        assert isinstance(self.forbidden_token_ids, list)

        self.cache = None
        self.context = ""
        self.generated_tokens = []  # track generated tokens
        self.generated_special_tokens = []  # track generated special tokens
        self.reset()
        self.embeds = None
        self.system_embeds = None

        # sliding window related states
        self._unit_history: List[Dict[str, Any]] = []
        self._next_unit_id: int = 0
        self._pending_unit_id: Optional[int] = None
        self._pending_unit_start_cache_len: int = 0
        self._system_preserve_length: int = 0
        self._position_offset: int = 0
        self._window_config = DuplexWindowConfig()
        self._window_enabled: bool = True
        self._rope_inv_freq_cache: Dict[Tuple, torch.Tensor] = {}

        # context preserving sliding window states
        # initial cache layout: [prefix] [suffix] [units...]
        # after first sliding window: [prefix] [previous_marker + content] [suffix] [units...]
        #                              fixed     dynamic sliding region      fixed
        self._preserve_prefix_length: int = 0  # original prefix length (fixed)
        self._previous_content_length: int = 0  # previous content length (dynamic, including marker)
        self._suffix_token_ids: List[int] = []  # suffix token ids (e.g. <|im_end|>)

        # previous marker (added dynamically after first sliding window)
        self._previous_marker: str = "\n\nprevious: "  # fixed prefix marker
        self._previous_marker_token_ids: List[int] = []  # marker token ids (initialized)
        self._has_previous: bool = False  # whether previous marker has been added

        # previous content
        self._previous_text: str = ""  # accumulated generated text (without marker)
        self._previous_token_ids: List[int] = []  # previous full token ids (including marker)

        # validation statistics
        self._sliding_event_count: int = 0  # sliding window trigger count
        self._total_dropped_tokens: int = 0  # total dropped token count
        self._total_dropped_units: int = 0  # total dropped unit count

    def sliding_embeds(self):
        # tmp = system_embeds
        # tmp +-》 embeds after 5s
        # reset
        # feed
        pass

    def reset(self):
        self.context = ""
        self.cache = None
        self.generated_tokens = []
        self.generated_special_tokens = []
        self.embeds = None
        self.system_embeds = None

        # sliding window state reset
        old_unit_count = len(self._unit_history) if hasattr(self, "_unit_history") else 0
        self._unit_history = []
        self._next_unit_id = 0
        self._pending_unit_id = None
        self._pending_unit_start_cache_len = 0
        self._system_preserve_length = 0
        self._position_offset = 0
        self._rope_inv_freq_cache = {}

        # context preserving sliding window state reset
        self._preserve_prefix_length = 0
        self._previous_content_length = 0
        self._suffix_token_ids = []
        self._previous_marker = "\n\nprevious: "
        self._previous_marker_token_ids = []
        self._has_previous = False
        self._previous_text = ""
        self._previous_token_ids = []

        # validation statistics
        self._sliding_event_count = 0  # sliding window trigger count
        self._total_dropped_tokens = 0  # total dropped token count
        self._total_dropped_units = 0  # total dropped unit count

    def get_cache_length(self) -> int:
        if self.cache is None:
            return 0
        if isinstance(self.cache, DynamicCache):
            if len(self.cache.key_cache) > 0 and self.cache.key_cache[0].numel() > 0:
                return self.cache.key_cache[0].shape[2]
            return 0
        # Tuple cache format
        return self.cache[0][0].shape[2]

    def get_total_generated_tokens(self) -> int:
        return sum(len(u.get("generated_tokens", [])) for u in self._unit_history)

    def register_unit_start(self) -> int:
        self._pending_unit_id = self._next_unit_id
        self._pending_unit_start_cache_len = self.get_cache_length()
        return self._pending_unit_id

    def register_unit_end(
        self,
        input_type: str,
        generated_tokens: Optional[List[int]] = None,
        is_listen: bool = False,
        generated_text: Optional[str] = None,
    ):
        """Call when unit ends, record unit information

        Should be called after feeding </unit> token

        Args:
            input_type: "audio" / "video" / "omni" / "system"
            generated_tokens: tokens generated by the unit (token ids)
            is_listen: whether the unit is in listen state
            generated_text: text generated by the unit (used for context preserving mode)
        """
        if self._pending_unit_id is None:
            logger.warning("register_unit_end called without register_unit_start")
            return

        # calculate the length of the unit
        current_cache_len = self.get_cache_length()
        unit_len = current_cache_len - self._pending_unit_start_cache_len

        if unit_len > 0:
            entry = {
                "unit_id": self._pending_unit_id,
                "length": unit_len,
                "type": input_type,
                "generated_tokens": generated_tokens or [],
                "generated_text": generated_text or "",  # used for context preserving mode
                "is_listen": is_listen,
            }
            self._unit_history.append(entry)

        self._pending_unit_id = None
        self._pending_unit_start_cache_len = 0
        self._next_unit_id += 1

    def register_system_prompt(self):
        """Call after system prompt prefill, record preserve length"""
        self._system_preserve_length = self.get_cache_length()

    # sliding window core methods

    def _get_rope_theta(self) -> float:
        """get model rope_theta configuration"""
        return float(getattr(self.m.config, "rope_theta", 10000.0))

    def _drop_tokens_from_cache(self, length: int) -> bool:
        """remove specified number of tokens from cache (protect system prompt)

        remove tokens in the range [preserve, preserve + length)
        supports DynamicCache and tuple cache formats
        """
        if self.cache is None or length <= 0:
            return False

        cache_type = "DynamicCache" if isinstance(self.cache, DynamicCache) else "TupleCache"
        cache_len_before = self.get_cache_length()
        offset_before = self._position_offset

        new_cache, new_offset, success = drop_tokens_from_cache(
            cache=self.cache,
            length=length,
            preserve=self._system_preserve_length,
            position_offset=self._position_offset,
            rope_theta=self._get_rope_theta(),
            inv_freq_cache=self._rope_inv_freq_cache,
        )
        if success:
            self.cache = new_cache  # For DynamicCache this is the same object (in-place)
            self._position_offset = new_offset

        return success

    def _drop_unit(self, unit_id: int) -> bool:
        """remove specified unit"""
        entries = [u for u in self._unit_history if u["unit_id"] == unit_id]
        if not entries:
            return False

        total_len = sum(e["length"] for e in entries)
        if total_len <= 0:
            for e in entries:
                self._unit_history.remove(e)
            return False

        if not self._drop_tokens_from_cache(total_len):
            return False

        for e in entries:
            self._unit_history.remove(e)

        return True

    def _drop_next_unit(self) -> bool:
        """remove the earliest non-system unit"""
        for entry in self._unit_history:
            unit_id = entry.get("unit_id")
            if unit_id is None:
                continue
            # skip system type
            if entry.get("type") == "system":
                continue
            if self._drop_unit(unit_id):
                return True
        return False

    def enforce_window(self) -> bool:
        """enforce sliding window strategy (same as single-mode, only look at cache length)

        when cache length exceeds high water line, loop to remove the earliest unit,
        until cache length drops below the low water line.
        """
        if not self._window_enabled:
            return False

        cfg = self._window_config
        cache_len_before = self.get_cache_length()

        if cache_len_before <= cfg.basic_window_high_tokens:
            return False  # not above high water line, no trigger

        dropped_count = 0
        cache_len = cache_len_before
        while cache_len > cfg.basic_window_low_tokens:
            if not self._drop_next_unit():
                break
            dropped_count += 1
            cache_len = self.get_cache_length()

        if dropped_count > 0:
            # update statistics counters
            self._sliding_event_count += 1
            self._total_dropped_tokens += cache_len_before - cache_len
            self._total_dropped_units += dropped_count

            # consistency check
            expected = self._system_preserve_length + sum(u["length"] for u in self._unit_history)
            is_consistent = expected == cache_len
            if not is_consistent:
                logger.error(
                    "CONSISTENCY ERROR! preserve=%d + sum(units)=%d != cache=%d, offset=%d",
                    self._system_preserve_length,
                    sum(u["length"] for u in self._unit_history),
                    cache_len,
                    self._position_offset,
                )

        return dropped_count > 0

    # context preserving sliding window methods

    def register_system_prompt_with_context(
        self,
        suffix_token_ids: Optional[List[int]] = None,
        context_previous_marker: str = "\n\nprevious: ",
    ):
        """register system prompt (with context preserving mode)

        initial cache layout: [prefix] [suffix] [units...]
        after first sliding window: [prefix] [context_previous_marker + content] [suffix] [units...]

        when calling this method, cache should only have prefix (without previous marker)
        suffix will be fed in later

        Args:
            suffix_token_ids: suffix token ids (e.g. id of <|im_end|>)
            context_previous_marker: previous marker prefix, e.g. "\\n\\nprevious: "
        """
        # prefix = current cache content (fixed, without previous marker)
        self._preserve_prefix_length = self.get_cache_length()
        self._previous_content_length = 0  # initially no previous content
        self._suffix_token_ids = suffix_token_ids or []
        # total preserve length = prefix + suffix (initially no previous)
        self._system_preserve_length = self._preserve_prefix_length + len(self._suffix_token_ids)

        # initialize previous related states
        self._previous_marker = context_previous_marker
        self._previous_marker_token_ids = (
            self.tokenizer.encode(context_previous_marker, add_special_tokens=False) if self.tokenizer else []
        )
        self._has_previous = False
        self._previous_text = ""
        self._previous_token_ids = []

    def _extract_generated_text(self, units: List[Dict[str, Any]]) -> Tuple[str, List[int]]:
        """extract generated text and token ids from units

        Args:
            units: list of units to extract

        Returns:
            (text, token_ids): concatenated text and token ids (filtered out special tokens)
        """
        text_parts = []
        token_ids = []

        for u in units:
            # only keep generated content of non-listen units
            if u.get("is_listen", False):
                continue
            gen_text = u.get("generated_text", "")
            gen_tokens = u.get("generated_tokens", [])

            # filter out special tokens from text
            if gen_text:
                clean_text = gen_text
                for st in self._all_special_tokens_text:
                    clean_text = clean_text.replace(st, "")
                if clean_text.strip():
                    text_parts.append(clean_text)

            # filter out special tokens
            if gen_tokens:
                filtered_tokens = [t for t in gen_tokens if t not in self._all_special_ids]
                token_ids.extend(filtered_tokens)

        return "".join(text_parts), token_ids

    def _rebuild_cache_with_previous(
        self,
        new_previous_tokens: List[int],
        units_to_keep_len: Optional[int] = None,
    ) -> bool:
        """rebuild cache, insert new previous content between prefix and suffix

        cache layout change:
        [prefix] [old_prev] [suffix] [old_units]  →  [prefix] [new_prev] [suffix] [remaining_units]

        Args:
            new_previous_tokens: new previous token ids
            units_to_keep_len: length of units to keep (from cache end backwards)
                                if None, calculate based on unit_history

        Returns:
            whether successful rebuild
        """
        if self.cache is None:
            return False

        old_previous_len = self._previous_content_length
        new_previous_len = len(new_previous_tokens)
        suffix_len = len(self._suffix_token_ids)
        total_cache_len = self.get_cache_length()

        # calculate length of units to keep
        if units_to_keep_len is None:
            units_to_keep_len = sum(u["length"] for u in self._unit_history)

        # special case: if previous is unchanged (new and old are empty), no need to rebuild prefix+suffix part of cache
        # but still need to reindex units RoPE (because a unit was deleted, position changed)
        if new_previous_len == 0 and old_previous_len == 0:
            # cache layout: [prefix(7)] [suffix(1)] [units...]
            # only keep prefix + suffix + remaining_units
            preserve_len = self._preserve_prefix_length + suffix_len

            # simply slice cache: [prefix+suffix] + [remaining_units]
            # remaining_units in cache end
            if units_to_keep_len > 0:
                # [0:preserve_len] + [total-units_to_keep_len:total]
                prefix_suffix_cache = self._slice_cache(0, preserve_len)
                units_cache = self._slice_cache(total_cache_len - units_to_keep_len, None)

                # calculate number of dropped tokens
                dropped_tokens = total_cache_len - preserve_len - units_to_keep_len

                # reindex units RoPE: position from (preserve_len + dropped_tokens) to preserve_len
                # note: no position_offset, because cache position has been compressed (from 0 start)
                if dropped_tokens > 0:
                    old_start = preserve_len + dropped_tokens
                    new_start = preserve_len
                    units_cache = self._reindex_rope_for_cache(units_cache, old_start, new_start, units_to_keep_len)

                self.cache = self._concat_caches(prefix_suffix_cache, units_cache)
            else:
                self.cache = self._slice_cache(0, preserve_len)

            return True

        # 1. get prefix cache (fixed)
        prefix_end = self._preserve_prefix_length
        prefix_cache = self._slice_cache(0, prefix_end)

        # 2. get units cache to keep (from end)
        units_start_in_old_cache = total_cache_len - units_to_keep_len
        units_cache = None
        if units_to_keep_len > 0:
            units_cache = self._slice_cache(units_start_in_old_cache, None)

        # 3. calculate new previous + suffix cache (needs forward)
        # merge previous tokens and suffix tokens
        prev_suffix_tokens = new_previous_tokens + self._suffix_token_ids
        prev_suffix_len = len(prev_suffix_tokens)

        new_prefix_prev_suffix_cache = prefix_cache
        if prev_suffix_len > 0:
            # Embed tokens
            prev_suffix_embeds = self.embed_tokens(prev_suffix_tokens)
            # calculate start position (after prefix)
            start_pos = self._preserve_prefix_length + self._position_offset

            # forward calculate KV cache
            with torch.no_grad():
                device = prev_suffix_embeds.device
                position_ids = torch.arange(
                    start_pos,
                    start_pos + prev_suffix_len,
                    device=device,
                ).unsqueeze(0)

                # use prefix cache as past_key_values
                outputs = self.m(
                    inputs_embeds=(
                        prev_suffix_embeds.unsqueeze(0) if prev_suffix_embeds.dim() == 2 else prev_suffix_embeds
                    ),
                    position_ids=position_ids,
                    past_key_values=prefix_cache,
                    use_cache=True,
                    return_dict=True,
                )
                # new cache contains prefix + new_previous + suffix
                new_prefix_prev_suffix_cache = outputs.past_key_values

        # 4. adjust units cache RoPE
        # new layout: [prefix] [new_prev] [suffix] [units]
        # note: no position_offset, because cache position has been compressed (from 0 start)
        new_system_total = prefix_end + new_previous_len + suffix_len
        if units_cache is not None and self._get_cache_len(units_cache) > 0:
            old_start = units_start_in_old_cache
            new_start = new_system_total

            if old_start != new_start:
                units_cache = self._reindex_rope_for_cache(units_cache, old_start, new_start, units_to_keep_len)

        # 5. concatenate new cache
        if units_cache is not None and self._get_cache_len(units_cache) > 0:
            self.cache = self._concat_caches(new_prefix_prev_suffix_cache, units_cache)
        else:
            self.cache = new_prefix_prev_suffix_cache

        # 6. update length
        self._previous_content_length = new_previous_len
        # total preserve length = prefix + previous + suffix
        self._system_preserve_length = prefix_end + new_previous_len + suffix_len

        # print detailed cache layout information
        prev_text_preview = self._previous_text[:50] + "..." if len(self._previous_text) > 50 else self._previous_text
        suffix_preview = self.tokenizer.decode(self._suffix_token_ids) if self._suffix_token_ids else ""
        return True

    def _slice_cache(self, start: int, end: Optional[int], clone: bool = True):
        """slice cache

        Args:
            start: start position
            end: end position (None means to end)
            clone: whether to clone (default True, to prevent shared memory issues)
        """
        if self.cache is None:
            return None
        if isinstance(self.cache, DynamicCache):
            # DynamicCache
            new_key_cache = [
                k[:, :, start:end, :].clone() if clone else k[:, :, start:end, :] for k in self.cache.key_cache
            ]
            new_value_cache = [
                v[:, :, start:end, :].clone() if clone else v[:, :, start:end, :] for v in self.cache.value_cache
            ]
            new_cache = DynamicCache()
            new_cache.key_cache = new_key_cache
            new_cache.value_cache = new_value_cache
            return new_cache
        else:
            # Tuple cache
            if clone:
                return tuple(
                    (layer[0][:, :, start:end, :].clone(), layer[1][:, :, start:end, :].clone()) for layer in self.cache
                )
            else:
                return tuple((layer[0][:, :, start:end, :], layer[1][:, :, start:end, :]) for layer in self.cache)

    @staticmethod
    def _get_cache_len(cache) -> int:
        if cache is None:
            return 0
        if isinstance(cache, DynamicCache):
            if len(cache.key_cache) > 0 and cache.key_cache[0].numel() > 0:
                return cache.key_cache[0].shape[2]
            return 0

        if cache and cache[0] and cache[0][0] is not None:
            return cache[0][0].shape[2]
        return 0

    @staticmethod
    def _concat_caches(cache1, cache2):
        if cache1 is None:
            return cache2
        if cache2 is None:
            return cache1

        if isinstance(cache1, DynamicCache):
            new_cache = DynamicCache()
            new_cache.key_cache = [torch.cat([k1, k2], dim=2) for k1, k2 in zip(cache1.key_cache, cache2.key_cache)]
            new_cache.value_cache = [
                torch.cat([v1, v2], dim=2) for v1, v2 in zip(cache1.value_cache, cache2.value_cache)
            ]
            return new_cache
        else:
            return tuple(
                (
                    torch.cat([layer1[0], layer2[0]], dim=2),
                    torch.cat([layer1[1], layer2[1]], dim=2),
                )
                for layer1, layer2 in zip(cache1, cache2)
            )

    def _reindex_rope_for_cache(self, cache, old_start: int, new_start: int, length: int):
        """reindex RoPE position for cache"""
        if cache is None or length <= 0:
            return cache

        if isinstance(cache, DynamicCache):
            device = cache.key_cache[0].device if cache.key_cache else None
        else:
            device = cache[0][0].device if cache and cache[0] else None

        if device is None:
            return cache

        old_positions = torch.arange(old_start, old_start + length, device=device, dtype=torch.long)
        new_positions = torch.arange(new_start, new_start + length, device=device, dtype=torch.long)

        rope_theta = self._get_rope_theta()

        if isinstance(cache, DynamicCache):
            new_key_cache = []
            for k in cache.key_cache:
                new_k = realign_rotary_suffix(k, old_positions, new_positions, rope_theta, self._rope_inv_freq_cache)
                new_key_cache.append(new_k)
            cache.key_cache = new_key_cache
            return cache
        else:
            new_cache = []
            for layer in cache:
                new_k = realign_rotary_suffix(
                    layer[0], old_positions, new_positions, rope_theta, self._rope_inv_freq_cache
                )
                new_cache.append((new_k, layer[1]))
            return tuple(new_cache)

    def _update_previous(
        self,
        new_text: str,
        new_tokens: List[int],
        max_tokens: int,
    ) -> None:
        """update previous context (also update cache)

        when first sliding window, dynamically add marker + text, subsequent sliding window append text
        when content exceeds max_tokens, truncate content (keep marker)
        rebuild cache to maintain consistency

        Args:
            new_text: new text
            new_tokens: new token ids
            max_tokens: previous content maximum token count (without marker)
        """
        marker_len = len(self._previous_marker_token_ids)
        tokens_to_drop = 0

        # if no new content, do not add marker, but still need to rebuild cache
        if not new_tokens and not new_text:
            # still need to rebuild cache (because a unit was deleted)
            self._rebuild_cache_with_previous(self._previous_token_ids)
            return

        if not self._has_previous:
            # when first has actual content: add marker + text
            self._previous_text = new_text
            self._previous_token_ids = self._previous_marker_token_ids.copy() + new_tokens
            self._has_previous = True
        else:
            # subsequent sliding window: append text to previous
            self._previous_text += new_text
            self._previous_token_ids.extend(new_tokens)

        # calculate token count of content (without marker)
        content_token_count = len(self._previous_token_ids) - marker_len

        # check if need to truncate content (keep marker)
        if content_token_count > max_tokens:
            # truncate left content, keep marker + latest max_tokens content
            tokens_to_drop = content_token_count - max_tokens
            old_text = self._previous_text
            # keep marker + truncated content
            content_tokens = self._previous_token_ids[marker_len + tokens_to_drop :]
            self._previous_token_ids = self._previous_marker_token_ids.copy() + content_tokens
            # redecode text (only decode content part)
            try:
                self._previous_text = self.tokenizer.decode(
                    content_tokens,
                    skip_special_tokens=True,
                )
            except Exception as e:
                logger.warning("_update_previous: decode failed: %s", e)

        # rebuild cache
        self._rebuild_cache_with_previous(self._previous_token_ids)

    def _drop_unit_with_context(
        self,
        unit_id: int,
        max_previous_tokens: int,
    ) -> Tuple[bool, str, List[int]]:
        """remove specified unit and return its generated content (for context preserving)

        process:
        1. extract generated content of unit
        2. remove unit from cache (without prefix+previous)
        3. append generated content to previous
        4. rebuild cache (in _update_previous)

        Args:
            unit_id: unit ID to remove
            max_previous_tokens: previous maximum token count

        Returns:
            (success, extracted_text, extracted_tokens): whether successful, extracted text and tokens
        """
        entries = [u for u in self._unit_history if u["unit_id"] == unit_id]
        if not entries:
            return False, "", []

        # extract generated content
        extracted_text, extracted_tokens = self._extract_generated_text(entries)

        # calculate total length
        total_len = sum(e["length"] for e in entries)
        if total_len <= 0:
            for e in entries:
                self._unit_history.remove(e)
            return False, extracted_text, extracted_tokens

        cache_before = self.get_cache_length()

        # remove from unit_history (record for later processing)
        for e in entries:
            self._unit_history.remove(e)

        # note: here no longer call _drop_tokens_from_cache
        # because _update_previous will rebuild the entire cache

        # update previous (also rebuild cache)
        self._update_previous(extracted_text, extracted_tokens, max_previous_tokens)

        return True, extracted_text, extracted_tokens

    def _drop_next_unit_with_context(self, max_previous_tokens: int) -> bool:
        """remove the earliest non-system unit (with context preserving)"""
        for entry in self._unit_history:
            unit_id = entry.get("unit_id")
            if unit_id is None:
                continue
            if entry.get("type") == "system":
                continue
            success, _, _ = self._drop_unit_with_context(unit_id, max_previous_tokens)
            if success:
                return True
        return False

    def enforce_window_with_context(self) -> bool:
        """context preserving sliding window execution

        when unit count exceeds max_units, remove the earliest unit,
        and accumulate its generated content to previous.
        Cache will be automatically rebuilt in _update_previous.

        Returns:
            whether sliding window is executed
        """
        if not self._window_enabled:
            return False

        cfg = self._window_config

        if cfg.sliding_window_mode != "context":
            # if not context mode, fallback to basic sliding window
            return self.enforce_window()

        cache_len_before = self.get_cache_length()
        units_before = len(self._unit_history)

        # context preserving mode: only check if unit count exceeds limit
        # (previous exceeds limit in _update_previous will automatically truncate left)
        if units_before <= cfg.context_max_units:
            return False

        # sliding window loop: remove unit until count ≤ max_units
        dropped_count = 0
        while len(self._unit_history) > cfg.context_max_units:
            if not self._drop_next_unit_with_context(cfg.context_previous_max_tokens):
                break

            dropped_count += 1

        cache_len_after = self.get_cache_length()

        if dropped_count > 0:
            # update statistics counter
            self._sliding_event_count += 1
            self._total_dropped_tokens += cache_len_before - cache_len_after
            self._total_dropped_units += dropped_count

            # consistency check
            expected = self._system_preserve_length + sum(u["length"] for u in self._unit_history)

        return dropped_count > 0

    def get_previous_context(self) -> Tuple[str, List[int]]:
        """get current accumulated previous context

        Returns:
            (previous_text, previous_token_ids): current accumulated text and token ids
        """
        return self._previous_text, self._previous_token_ids.copy()

    def get_window_stats(self) -> Dict[str, Any]:
        """get sliding window statistics"""
        unit_lengths = [u["length"] for u in self._unit_history]
        return {
            "cache_length": self.get_cache_length(),
            "unit_count": len(self._unit_history),
            "unit_lengths": unit_lengths,
            "unit_total_length": sum(unit_lengths),
            "system_preserve_length": self._system_preserve_length,
            "position_offset": self._position_offset,
            "window_enabled": self._window_enabled,
            "total_generated_tokens": self.get_total_generated_tokens(),
            "pending_unit_id": self._pending_unit_id,
            "next_unit_id": self._next_unit_id,
            "config": {
                "sliding_window_mode": self._window_config.sliding_window_mode,
                "basic_window_high_tokens": self._window_config.basic_window_high_tokens,
                "basic_window_low_tokens": self._window_config.basic_window_low_tokens,
                "context_previous_max_tokens": self._window_config.context_previous_max_tokens,
                "context_max_units": self._window_config.context_max_units,
            },
            # context preserving related
            "preserve_prefix_length": self._preserve_prefix_length,
            "previous_content_length": self._previous_content_length,
            "suffix_token_count": len(self._suffix_token_ids),
            "previous_text_length": len(self._previous_text),
            "previous_token_count": len(self._previous_token_ids),
            "has_system_template": self._system_prompt_template is not None,
        }

    def _verify_consistency(self) -> bool:
        """verify unit history and cache length consistency"""
        expected = self._system_preserve_length + sum(u["length"] for u in self._unit_history)
        actual = self.get_cache_length()
        return expected == actual

    def print_verification_summary(self) -> Dict[str, Any]:
        """print verification summary (for comparing off/basic/context mode)

        Returns:
            dictionary containing key verification data
        """
        cfg = self._window_config

        # collect all generated text
        all_generated_text = []
        all_generated_tokens = []
        for u in self._unit_history:
            if not u.get("is_listen", False):
                gen_text = u.get("generated_text", "")
                gen_tokens = u.get("generated_tokens", [])
                if gen_text:
                    all_generated_text.append(gen_text)
                if gen_tokens:
                    all_generated_tokens.extend(gen_tokens)

        combined_text = "".join(all_generated_text)

        summary = {
            "mode": cfg.sliding_window_mode,
            "final_cache_length": self.get_cache_length(),
            "final_unit_count": len(self._unit_history),
            "sliding_event_count": self._sliding_event_count,
            "total_dropped_tokens": self._total_dropped_tokens,
            "total_dropped_units": self._total_dropped_units,
            "total_generated_tokens": len(all_generated_tokens),
            "generated_text": combined_text,
            "previous_text": self._previous_text,
            "previous_token_count": len(self._previous_token_ids),
            "position_offset": self._position_offset,
            "system_preserve_length": self._system_preserve_length,
        }

        return summary

    def set_window_config(self, config: DuplexWindowConfig) -> None:
        """set sliding window configuration"""
        self._window_config = config

    def set_window_enabled(self, enabled: bool) -> None:
        """enable/disable sliding window"""
        old_enabled = self._window_enabled
        self._window_enabled = enabled

    def get_context(self):
        return self.context

    def embed_token(self, tid):
        if isinstance(tid, int):
            tid = torch.tensor([tid], device=self.m.device)
        return self.m.model.embed_tokens(tid)

    def embed_tokens(self, token_ids: List[int]) -> torch.Tensor:
        """batch embed multiple tokens

        Args:
            token_ids: list of token ids

        Returns:
            embeddings tensor [L, H]
        """
        if not token_ids:
            return torch.empty(0, self.m.config.hidden_size, device=self.m.device)
        tids = torch.tensor(token_ids, device=self.m.device)
        return self.m.model.embed_tokens(tids)

    @torch.no_grad()
    def feed(self, embeds: torch.Tensor, return_logits: bool = False):
        """
        embeds : [L, H]   —— new embedding sequence fed into model at once
        """
        L = embeds.size(0)
        device = embeds.device

        past_len = self.get_cache_length()
        pos_ids = torch.arange(past_len, past_len + L, device=device).unsqueeze(0)  # [1, L]

        out = self.m(
            inputs_embeds=embeds.unsqueeze(0),  # [1, L, H]
            position_ids=pos_ids,
            past_key_values=self.cache,
            # use_cache = True,
            return_dict=True,
            output_hidden_states=True,
            # attention_mask=attention_mask
        )
        self.cache = out.past_key_values

        if return_logits:
            logits = self.m.lm_head(out.hidden_states[-1])[:, -1]  # [1, vocab]
            return logits, out.hidden_states[-1]

    @torch.no_grad()
    def decode(
        self,
        logits,
        mode: Literal["sampling", "greedy"] = "sampling",
        temperature=0.7,
        top_k=20,
        top_p=0.8,
        listen_top_k=None,
        listen_prob_scale=1.0,
        text_repetition_penalty=1.05,
        text_repetition_window_size=512,
    ):
        """
        Args:
            logits:
            mode: sampling or greedy
            temperature:
            top_k:
            top_p:
            listen_top_k: force listen_id to be in top-k to keep
            listen_prob_scale: multiply listen_id probability by a weight (<1 means decrease, >1 means increase)
            text_repetition_penalty: repetition penalty coefficient, >1.0 means decrease repetition, <1.0 means increase repetition
            text_repetition_window_size: repetition penalty window size

        Sampling strategy:
            1. first sample all tokens with original logits (apply temperature)
            2. if sampled chunk_eos, return directly (keep the original model's decision of when to stop)
            3. if not sampled chunk_eos, mask it (set logit to -inf), continue sampling text tokens
            4. apply repetition penalty, top-k, top-p, etc. to the text tokens for the final sampling
        """

        logits = logits.clone()

        # 0. independently check chunk_eos before sampling
        eos_id = self.chunk_eos_id

        with torch.no_grad():
            if mode == "greedy":
                sampled_token = torch.argmax(logits[0]).item()
            else:
                original_probs = F.softmax(logits[0], dim=-1)
                sampled_token = torch.multinomial(original_probs, num_samples=1).item()

            # if sampled chunk_eos, return directly
            if sampled_token == eos_id:
                next_token_id = torch.tensor([eos_id], device=logits.device)
                next_token_str = self.tokenizer.decode(next_token_id)

                return next_token_id

        # if not sampled chunk_eos, set its logit to -inf
        if self.forbidden_token_ids:
            logits[:, self.forbidden_token_ids] = float("-inf")

        # 1. apply repetition penalty
        if text_repetition_penalty != 1.0 and len(self.generated_tokens) > 0:
            # get recent tokens (within window size) considering special tokens and normal tokens
            recent_tokens = self.generated_tokens[-text_repetition_window_size:]

            # make it unique
            recent_tokens = list(set(recent_tokens))

            # apply penalty to repeated tokens
            for token_id in recent_tokens:
                if token_id < logits.size(-1):  # ensure token_id is in vocabulary range
                    if text_repetition_penalty > 1.0:
                        # penalize repetition: decrease logits
                        logits[0, token_id] /= text_repetition_penalty
                    else:
                        # encourage repetition: increase logits
                        logits[0, token_id] *= 1.0 / text_repetition_penalty

        if listen_prob_scale != 1.0:  # modify listen token logit separately
            logits[0, self.listen_id] *= listen_prob_scale

        listen_rank = (logits[0] > logits[0, self.listen_id]).sum().item()

        if listen_top_k is not None and listen_rank < listen_top_k:  # listen_id is in top-k, return directly
            next_token_id = torch.tensor([self.listen_id], device=logits.device)
            next_token_str = self.tokenizer.decode(next_token_id)

            if next_token_str == "<|listen|>":
                self.context += " "
            else:
                self.context += next_token_str

            return next_token_id

        if mode == "greedy":
            next_token_id = torch.argmax(logits, dim=-1)
        elif mode == "sampling":
            logits = logits / temperature
            logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
            probs = F.softmax(logits, dim=-1)
            next_token_id = torch.multinomial(probs, num_samples=1).squeeze(1)
        else:
            raise ValueError(f"Unsupported decode mode: {mode}")

        if next_token_id.item() not in self.special_token_ids:
            self.generated_tokens.append(next_token_id.item())
        else:
            self.generated_special_tokens.append(next_token_id.item())

        return next_token_id


def _download_url_to_tempfile(url: str, suffix: str = "", timeout: int = 60) -> str:
    """
    Download a URL to a temporary file and return the path.

    Args:
        url: HTTP/HTTPS URL to download
        suffix: File suffix (e.g., ".jpg", ".wav", ".mp4")
        timeout: Download timeout in seconds

    Returns:
        Path to the downloaded temporary file
    """
    import tempfile

    import requests

    response = requests.get(url, timeout=timeout)
    response.raise_for_status()

    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
        f.write(response.content)
        return f.name


def _is_url(path: str) -> bool:
    return path.startswith(("http://", "https://"))


def normalize_content_item(item) -> Union[str, Any, List[Any]]:
    """Normalize structured content item to native format.

    Supports:
    - Native format: str, PIL.Image, np.ndarray (pass through)
    - OpenAI structured format:
        - {"type": "text", "text": "..."} -> str
        - {"type": "image_url", "image_url": {"url": "..."}} -> PIL.Image
        - {"type": "audio_url", "audio_url": {"url": "..."}} -> np.ndarray
        - {"type": "video_url", "video_url": {"url": "...", ...}} -> List[Image, ndarray, ...]

    URL formats supported:
        - Local file path: "/path/to/file.jpg"
        - HTTP/HTTPS URL: "https://example.com/image.jpg"

    Args:
        item: Content item to normalize

    Returns:
        Normalized item. For video_url, returns a tuple ("__video_contents__", list)
        that will be flattened by normalize_content().

    Raises:
        ValueError: If content type is unknown or unsupported
    """
    import os

    import numpy as np
    from PIL import Image

    if isinstance(item, str):
        return item
    if isinstance(item, Image.Image):
        return item
    if isinstance(item, np.ndarray):
        return item

    if isinstance(item, dict):
        item_type = item.get("type")

        if item_type == "text":
            return item.get("text", "")

        elif item_type == "image_url":
            image_url_obj = item.get("image_url", {})
            url = image_url_obj.get("url", "") if isinstance(image_url_obj, dict) else image_url_obj

            if _is_url(url):
                # Download to temp file
                temp_path = _download_url_to_tempfile(url, suffix=".jpg", timeout=30)
                img = Image.open(temp_path)
                os.unlink(temp_path)
                return img
            else:
                return Image.open(url)
        elif item_type == "audio_url":
            import librosa

            audio_url_obj = item.get("audio_url", {})
            url = audio_url_obj.get("url", "") if isinstance(audio_url_obj, dict) else audio_url_obj

            if _is_url(url):
                # Download to temp file
                temp_path = _download_url_to_tempfile(url, suffix=".wav", timeout=60)
                audio_np, _ = librosa.load(temp_path, sr=16000, mono=True)
                os.unlink(temp_path)
                return audio_np
            else:
                audio_np, _ = librosa.load(url, sr=16000, mono=True)
                return audio_np
        elif item_type == "video_url":
            # Video processing - returns a LIST of items (frames + audio segments)
            # Note: Unlike image_url/audio_url which return single items,
            # video_url returns a list that will be flattened into the content
            from minicpmo.utils import get_video_frame_audio_segments

            video_url_obj = item.get("video_url", {})
            if isinstance(video_url_obj, dict):
                video_url = video_url_obj.get("url", "")
                # Get optional parameters from video_url object (OpenAI style)
                stack_frames = video_url_obj.get("stack_frames", 1)
                use_ffmpeg = video_url_obj.get("use_ffmpeg", False)
                use_audio = video_url_obj.get("use_audio", True)
            else:
                video_url = video_url_obj
                stack_frames = 1
                use_ffmpeg = False
                use_audio = True

            # Handle HTTP/HTTPS URL - download to temp file
            temp_video_path = None
            if _is_url(video_url):
                temp_video_path = _download_url_to_tempfile(video_url, suffix=".mp4", timeout=120)
                video_path = temp_video_path
            else:
                video_path = video_url

            # Extract frames and audio segments
            video_frames, audio_segments, stacked_frames = get_video_frame_audio_segments(
                video_path,
                stack_frames=stack_frames,
                use_ffmpeg=use_ffmpeg,
            )

            # Clean up temp file if downloaded
            if temp_video_path is not None:
                os.unlink(temp_video_path)

            # Build omni_contents (interleaved frames and audio, or frames only)
            omni_contents = []
            for i in range(len(video_frames)):
                omni_contents.append(video_frames[i])
                if use_audio:
                    omni_contents.append(audio_segments[i])
                if stacked_frames is not None and i < len(stacked_frames) and stacked_frames[i] is not None:
                    omni_contents.append(stacked_frames[i])

            # Return as a special marker to be flattened later
            return "__video_contents__", omni_contents
        else:
            raise ValueError(f"Unknown content type: {item_type}")

    raise ValueError(f"Cannot normalize content item of type: {type(item)}")


def normalize_content(content) -> list:
    """Normalize message content to list of native items.

    Input formats:
    - str: "hello" -> ["hello"]
    - list of native items: [str, Image, np.ndarray] -> pass through with normalization
    - list of structured items: [{"type": "text", ...}] -> normalize each
    - video type: automatically expanded to omni_contents
    - mixed: works too

    Args:
        content: Message content in any supported format

    Returns:
        List of native items (str, PIL.Image, np.ndarray)

    Examples:
        >>> normalize_content("hello")
        ["hello"]

        >>> normalize_content([{"type": "text", "text": "hi"}])
        ["hi"]

        >>> normalize_content([{"type": "video", "video": "/path/to/video.mp4"}])
        [<PIL.Image>, <np.ndarray>, <PIL.Image>, <np.ndarray>, ...]
    """
    import numpy as np
    from PIL import Image

    if isinstance(content, str):
        return [content]

    if isinstance(content, list):
        result = []
        for item in content:
            normalized = normalize_content_item(item)
            # Handle video content (returns tuple with marker)
            if isinstance(normalized, tuple) and len(normalized) == 2 and normalized[0] == "__video_contents__":
                # Flatten video contents into result
                result.extend(normalized[1])
            else:
                result.append(normalized)
        return result

    # Single non-list item (Image or np.ndarray)
    if isinstance(content, (Image.Image, np.ndarray)):
        return [content]

    normalized = normalize_content_item(content)
    if isinstance(normalized, tuple) and len(normalized) == 2 and normalized[0] == "__video_contents__":
        return normalized[1]
    return [normalized]