Sy03andHongsheng Liu
0feb92bdfa
Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
93 lines
2.9 KiB
Python
93 lines
2.9 KiB
Python
from typing import Any
|
|
|
|
import torch
|
|
|
|
from vllm_omni.outputs.mm_outputs import MultimodalPayload
|
|
from vllm_omni.outputs.output_modality import DRAINABLE_MODALITIES
|
|
|
|
CHUNK_METADATA_KEYS: frozenset[str] = frozenset(
|
|
{
|
|
"audio_text_total_chars",
|
|
"duplex_epoch",
|
|
"duplex_turn_id",
|
|
"llm_output_text_utf8",
|
|
"segment_end",
|
|
"tts_is_last_chunk",
|
|
"turn_end",
|
|
}
|
|
)
|
|
|
|
|
|
def _is_chunk_metadata_key(key: str) -> bool:
|
|
if key.startswith("meta."):
|
|
return key.split(".", 1)[1] in CHUNK_METADATA_KEYS
|
|
return key in CHUNK_METADATA_KEYS
|
|
|
|
|
|
def replace_snapshot_keys(
|
|
accumulated: MultimodalPayload,
|
|
incoming: MultimodalPayload,
|
|
) -> None:
|
|
"""Replace per-chunk metadata with the latest values."""
|
|
for key in (*incoming.tensors, *incoming.metadata):
|
|
if _is_chunk_metadata_key(key):
|
|
accumulated.tensors.pop(key, None)
|
|
accumulated.metadata.pop(key, None)
|
|
|
|
|
|
def drain_delta_payload(payload: MultimodalPayload) -> None:
|
|
"""Remove client-facing delta data while retaining request-level state."""
|
|
for modality_key in DRAINABLE_MODALITIES:
|
|
key = str(modality_key)
|
|
payload.tensors.pop(key, None)
|
|
payload.metadata.pop(key, None)
|
|
|
|
for metadata_key in CHUNK_METADATA_KEYS:
|
|
flat_key = f"meta.{metadata_key}"
|
|
payload.tensors.pop(flat_key, None)
|
|
payload.metadata.pop(flat_key, None)
|
|
|
|
meta = payload.metadata.get("meta")
|
|
if isinstance(meta, dict):
|
|
filtered = {key: value for key, value in meta.items() if key not in CHUNK_METADATA_KEYS}
|
|
if filtered:
|
|
payload.metadata["meta"] = filtered
|
|
else:
|
|
payload.metadata.pop("meta", None)
|
|
|
|
|
|
def _payload_meta_value(payload: MultimodalPayload, key: str) -> Any:
|
|
flat_key = f"meta.{key}"
|
|
if flat_key in payload.tensors:
|
|
return payload.tensors[flat_key]
|
|
if flat_key in payload.metadata:
|
|
return payload.metadata[flat_key]
|
|
meta = payload.metadata.get("meta")
|
|
if isinstance(meta, dict):
|
|
return meta.get(key)
|
|
return None
|
|
|
|
|
|
def _last_scalar_int(value: Any) -> int | None:
|
|
if value is None:
|
|
return None
|
|
if isinstance(value, bool | int):
|
|
return int(value)
|
|
if isinstance(value, list | tuple):
|
|
return _last_scalar_int(value[-1]) if value else None
|
|
if isinstance(value, torch.Tensor):
|
|
try:
|
|
if value.numel() == 0:
|
|
return None
|
|
return int(value.detach().cpu().reshape(-1)[-1].item())
|
|
except (RuntimeError, TypeError, ValueError):
|
|
return None
|
|
return None
|
|
|
|
|
|
def is_non_final_delta_audio_chunk(payload: MultimodalPayload, mm_type: str | None) -> bool:
|
|
"""Return whether an audio delta explicitly declares more chunks."""
|
|
if str(mm_type or "").lower() != "audio" and "audio" not in payload:
|
|
return False
|
|
return _last_scalar_int(_payload_meta_value(payload, "tts_is_last_chunk")) == 0
|