Spaces:
Sleeping
Sleeping
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 (must precede torch) | |
| import json # noqa: E402 | |
| import re # noqa: E402 | |
| import time # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import librosa # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| from huggingface_hub import snapshot_download # noqa: E402 | |
| from transformers import ( # noqa: E402 | |
| AutoModelForCausalLM, | |
| AutoTokenizer, | |
| LogitsProcessor, | |
| LogitsProcessorList, | |
| WhisperFeatureExtractor, | |
| ) | |
| from hy_simulmt.model import build_llm_view, encode_audio, load_audio_tower, speech_token_count # noqa: E402 | |
| MODEL_ID = "tencent/Hy-SimulMT-30B-A3B" | |
| DTYPE = torch.bfloat16 | |
| SR = 16000 | |
| MAX_SECONDS = 60 | |
| # ---------------------------------------------------------------------------- | |
| # Load (module scope, eager .to("cuda") -- ZeroGPU streams weights in on first call) | |
| # ---------------------------------------------------------------------------- | |
| root = Path( | |
| snapshot_download(MODEL_ID, allow_patterns=["*.safetensors", "*.json", "*.jinja"]) | |
| ) | |
| view = build_llm_view(root, Path("/tmp/hy_simulmt_llm_view")) | |
| tokenizer = AutoTokenizer.from_pretrained(view) | |
| feature_extractor = WhisperFeatureExtractor.from_pretrained(root) | |
| llm = AutoModelForCausalLM.from_pretrained(view, dtype=DTYPE, attn_implementation="sdpa").eval() | |
| encoder, adapter = load_audio_tower(root, DTYPE) | |
| llm.to("cuda") | |
| encoder.to("cuda") | |
| adapter.to("cuda") | |
| print("Hy-SimulMT loaded", flush=True) | |
| # ---------------------------------------------------------------------------- | |
| # Protocol constants (SimulInfer/simulinfer/translation.py + templates/simulst.jinja) | |
| # ---------------------------------------------------------------------------- | |
| HEADER = "<｜hy_begin▁of▁sentence｜><｜reasoning_mode｜>reasoning_effort:no_think" | |
| USER, ASSISTANT, EOT = "<｜hy_User｜>", "<｜hy_Assistant｜><think></think>", "<｜hy_EOT｜>" | |
| AUDIO_BEGIN = tokenizer.convert_tokens_to_ids("<｜hy_place▁holder▁no▁6｜>") | |
| AUDIO_TOKEN = tokenizer.convert_tokens_to_ids("<｜hy_place▁holder▁no▁5｜>") | |
| AUDIO_END = tokenizer.convert_tokens_to_ids("<｜hy_place▁holder▁no▁8｜>") | |
| STOP_IDS = [ | |
| tokenizer.convert_tokens_to_ids(EOT), | |
| tokenizer.convert_tokens_to_ids("<｜hy_end▁of▁sentence｜>"), | |
| tokenizer.eos_token_id, | |
| ] | |
| INSTRUCTIONS = { | |
| "zh-en": "请完成下面的同声传译任务，将陆续收到的语音翻译成英文。请结合上下文边听边译，输出能够确定的译文，不必等待整句结束。", | |
| "en-zh": "请完成下面的同声传译任务，将陆续收到的语音翻译成中文。请结合上下文边听边译，输出能够确定的译文，不必等待整句结束。", | |
| } | |
| ST_INSTRUCTIONS = { | |
| "zh-en": "把这段语音翻译成英文，要全部输出，不要分块。", | |
| "en-zh": "把这段语音翻译成中文，要全部输出，不要分块。", | |
| } | |
| DIRECTIONS = {"Chinese → English": "zh-en", "English → Chinese": "en-zh"} | |
| CHUNK = 1600 * 16 # first_chunk_ms = chunk_ms = 1600 | |
| MAX_TOKENS, ST_MAX_TOKENS = 64, 256 | |
| ST_WAIT_CHUNKS = 4 | |
| SILENCE_RMS = 0.003 | |
| TRAILING_SILENCE = int(2.4 * SR) | |
| # Target-language mask for zh->en (SimulInfer target_language.py: block Han / partial-UTF8 tokens) | |
| def _blocked_han_ids(): | |
| han = re.compile("[⺀-〇㇀-㇯㐀-䶿一-鿿豈-\U00020000-\U0003347f]") | |
| values = list(range(33, 127)) + list(range(161, 173)) + list(range(174, 256)) | |
| chars, extra = values[:], 0 | |
| for v in range(256): | |
| if v not in values: | |
| values.append(v) | |
| chars.append(256 + extra) | |
| extra += 1 | |
| decode = dict(zip(map(chr, chars), values)) | |
| data = json.loads((root / "tokenizer.json").read_text()) | |
| blocked = set() | |
| for tok, tid in data["model"]["vocab"].items(): | |
| try: | |
| text = bytes(decode[c] for c in tok).decode("utf-8", errors="strict") | |
| except (KeyError, UnicodeDecodeError): | |
| blocked.add(tid) | |
| continue | |
| if han.search(text): | |
| blocked.add(tid) | |
| for item in data.get("added_tokens", []): | |
| if not item.get("special") and han.search(item["content"]): | |
| blocked.add(item["id"]) | |
| return sorted(blocked) | |
| HAN_IDS = torch.tensor(_blocked_han_ids(), dtype=torch.long) | |
| class BlockTokens(LogitsProcessor): | |
| def __init__(self, ids): | |
| self.ids = ids | |
| def __call__(self, input_ids, scores): | |
| scores[:, self.ids.to(scores.device)] = -float("inf") | |
| return scores | |
| # ---------------------------------------------------------------------------- | |
| # Inference helpers | |
| # ---------------------------------------------------------------------------- | |
| def pcm16(x: np.ndarray) -> np.ndarray: | |
| """SimulInfer sends model audio as PCM16 WAV; reproduce that quantisation.""" | |
| return (np.clip(np.rint(np.clip(x, -1, 1) * 32768.0), -32768, 32767) / 32768.0).astype(np.float32) | |
| def encode_part(samples: np.ndarray) -> torch.Tensor: | |
| f = feature_extractor([pcm16(samples)], sampling_rate=SR, return_attention_mask=True, | |
| padding="longest", return_tensors="pt") | |
| n = int(f["attention_mask"].sum()) | |
| emb = encode_audio(encoder, adapter, f["input_features"][0], n) | |
| assert emb.shape[0] == speech_token_count(n) | |
| return emb | |
| def render(messages): | |
| """messages: list of ('user', [str | Tensor ...]) / ('assistant', str). Returns ids + audio embeds.""" | |
| ids, audio = [], [] | |
| def text(s): | |
| ids.extend(tokenizer.encode(s, add_special_tokens=False)) | |
| text(HEADER) | |
| for role, content in messages: | |
| if role == "user": | |
| text(USER) | |
| for p in content: | |
| if isinstance(p, str): | |
| text(p) | |
| else: | |
| ids.append(AUDIO_BEGIN) | |
| ids.extend([AUDIO_TOKEN] * p.shape[0]) | |
| ids.append(AUDIO_END) | |
| audio.append(p) | |
| else: | |
| text(ASSISTANT + content + EOT) | |
| text(ASSISTANT) | |
| return ids, audio | |
| def generate(messages, max_new_tokens, direction, repetition_penalty=1.0): | |
| ids, audio = render(messages) | |
| input_ids = torch.tensor([ids], device="cuda") | |
| embeds = llm.get_input_embeddings()(input_ids) | |
| if audio: | |
| mask = input_ids[0] == AUDIO_TOKEN | |
| embeds[0, mask] = torch.cat(audio, 0).to(embeds.dtype) | |
| procs = LogitsProcessorList([BlockTokens(HAN_IDS)] if direction == "zh-en" else []) | |
| out = llm.generate( | |
| inputs_embeds=embeds, | |
| attention_mask=torch.ones_like(input_ids), | |
| max_new_tokens=max_new_tokens, | |
| do_sample=False, | |
| repetition_penalty=repetition_penalty, | |
| no_repeat_ngram_size=3, | |
| eos_token_id=STOP_IDS, | |
| pad_token_id=tokenizer.pad_token_id, | |
| logits_processor=procs, | |
| )[0].tolist() | |
| truncated = not (out and out[-1] in STOP_IDS) | |
| while out and out[-1] in STOP_IDS: | |
| out = out[:-1] | |
| return tokenizer.decode(out, skip_special_tokens=True), truncated | |
| def load_audio(path): | |
| if path is None: | |
| raise gr.Error("Please upload or record an audio clip.") | |
| audio, _ = librosa.load(path, sr=SR, mono=True) | |
| if not len(audio): | |
| raise gr.Error("The audio clip is empty.") | |
| return audio[: MAX_SECONDS * SR].astype(np.float32) | |
| def fmt_t(samples): | |
| return f"{samples / SR:5.1f}s" | |
| def _duration(audio_path, direction="Chinese → English", mode="Simultaneous", *args, **kwargs): | |
| try: | |
| secs = min(MAX_SECONDS, librosa.get_duration(path=audio_path)) | |
| except Exception: | |
| secs = MAX_SECONDS | |
| return int(8 + secs * (0.4 if mode == "Simultaneous" else 0.2)) | |
| def translate(audio_path: str, direction: str = "Chinese → English", mode: str = "Simultaneous"): | |
| """Translate Chinese/English speech with Hy-SimulMT. | |
| Args: | |
| audio_path: path or URL of a speech recording (Chinese or English, up to 60 s). | |
| direction: "Chinese → English" or "English → Chinese". | |
| mode: "Simultaneous" replays the clip in 1.6 s chunks like a live interpreter | |
| (incremental output); "Offline" translates the whole clip at once. | |
| Returns: | |
| The translated text and a per-chunk timeline (markdown). | |
| """ | |
| d = DIRECTIONS.get(direction, direction) | |
| if d not in INSTRUCTIONS: | |
| raise gr.Error("Unknown direction") | |
| audio = load_audio(audio_path) | |
| t0 = time.perf_counter() | |
| if mode != "Simultaneous": | |
| parts = [ST_INSTRUCTIONS[d]] + [encode_part(audio[i:i + 480000]) for i in range(0, len(audio), 480000)] | |
| text, _ = generate([("user", parts)], ST_MAX_TOKENS, d, repetition_penalty=1.05) | |
| yield text, f"Offline translation of {len(audio) / SR:.1f}s audio in {time.perf_counter() - t0:.1f}s." | |
| return | |
| messages, pending, rows = [], [], [] | |
| full = "" | |
| voiced_waits = continuous_silence = 0 | |
| st_count = 0 | |
| def st_flush(start_label): | |
| nonlocal full, voiced_waits, st_count | |
| if not any(not b[2] for b in pending): | |
| return "" | |
| samples = np.concatenate([b[0][: b[1]] for b in pending]) | |
| parts = [ST_INSTRUCTIONS[d]] + [encode_part(samples[i:i + 480000]) for i in range(0, len(samples), 480000)] | |
| text, _ = generate([("user", parts)], ST_MAX_TOKENS, d, repetition_penalty=1.05) | |
| st_count += 1 | |
| if not text.strip(): | |
| return "" | |
| full += text | |
| if messages and messages[-1][0] == "assistant": | |
| messages[-1] = ("assistant", text) | |
| pending.clear() | |
| voiced_waits = 0 | |
| return text | |
| def table(): | |
| head = "| Audio | Action | Output |\n|---|---|---|\n" | |
| return head + "\n".join(rows[-200:]) | |
| cursor = 0 | |
| n_chunks = (len(audio) + CHUNK - 1) // CHUNK | |
| for idx in range(n_chunks): | |
| chunk = audio[idx * CHUNK:(idx + 1) * CHUNK] | |
| valid = len(chunk) | |
| final = idx == n_chunks - 1 | |
| if valid < CHUNK: | |
| chunk = np.pad(chunk, (0, CHUNK - valid)) | |
| begin, cursor = cursor, cursor + valid | |
| label = f"{fmt_t(begin)}–{fmt_t(cursor)}" | |
| silent = float(np.sqrt(np.mean(chunk[:valid] ** 2))) < SILENCE_RMS | |
| continuous_silence = continuous_silence + valid if silent else 0 | |
| if silent and not messages: # skip_initial_silence | |
| rows.append(f"| {label} | skip (silence) | |") | |
| yield full, table() | |
| continue | |
| turn = ("user", ([INSTRUCTIONS[d]] if not messages else []) + [encode_part(chunk)]) | |
| text, truncated = generate(messages + [turn], MAX_TOKENS, d) | |
| messages += [turn, ("assistant", text)] | |
| pending.append((chunk, valid, silent)) | |
| if text.strip(): | |
| full += text | |
| pending.clear() | |
| voiced_waits = 0 | |
| rows.append(f"| {label} | EMIT | {text.strip()} |") | |
| else: | |
| if not silent: | |
| voiced_waits += 1 | |
| rows.append(f"| {label} | WAIT | |") | |
| periodic = voiced_waits >= ST_WAIT_CHUNKS | |
| silence_end = voiced_waits and silent and continuous_silence >= TRAILING_SILENCE | |
| if periodic or silence_end or final: | |
| st = st_flush(label) | |
| if st: | |
| rows.append(f"| {label} | ST recovery | {st.strip()} |") | |
| if truncated and not final: # max_tokens reached -> cut the run | |
| if any(not b[2] for b in pending): | |
| st = st_flush(label) | |
| if st: | |
| rows.append(f"| {label} | ST recovery | {st.strip()} |") | |
| messages = [] | |
| pending.clear() | |
| voiced_waits = 0 | |
| yield full, table() | |
| elapsed = time.perf_counter() - t0 | |
| rows.append(f"| done | {n_chunks} chunks, {st_count} ST recoveries | processed {len(audio) / SR:.1f}s audio in {elapsed:.1f}s |") | |
| yield full, table() | |
| # ---------------------------------------------------------------------------- | |
| # UI | |
| # ---------------------------------------------------------------------------- | |
| CSS = """ | |
| #col-container { max-width: 1100px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks() as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# Hy-SimulMT 30B-A3B\n" | |
| "Tencent Hunyuan's simultaneous speech translation model (Chinese ⇄ English). " | |
| "In **Simultaneous** mode the clip is replayed in 1.6 s chunks, like a live interpreter: " | |
| "for each chunk the model either emits translated text or waits for more context.\n\n" | |
| "[Model](https://huggingface.co/tencent/Hy-SimulMT-30B-A3B) · " | |
| "[GitHub](https://github.com/Tencent-Hunyuan/Hy-SimulMT) · " | |
| "Text output only; the official speech output and voice cloning need a separate TTS service." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio_in = gr.Audio(label="Source speech (max 60 s)", sources=["upload", "microphone"], type="filepath") | |
| direction = gr.Radio(list(DIRECTIONS), value="Chinese → English", label="Direction") | |
| mode = gr.Radio(["Simultaneous", "Offline"], value="Simultaneous", label="Mode", | |
| info="Offline translates the whole clip in one pass") | |
| run = gr.Button("Translate", variant="primary") | |
| with gr.Column(): | |
| out_text = gr.Textbox(label="Translation", lines=8) | |
| with gr.Accordion("Chunk timeline", open=True): | |
| timeline = gr.Markdown() | |
| gr.Examples( | |
| examples=[ | |
| ["examples/zh_lunar_south_pole.wav", "Chinese → English"], | |
| ["examples/en_reasonable_measures.wav", "English → Chinese"], | |
| ["examples/zh_womens_basketball.wav", "Chinese → English"], | |
| ["examples/en_learning_culture.wav", "English → Chinese"], | |
| ], | |
| inputs=[audio_in, direction], | |
| outputs=[out_text, timeline], | |
| fn=translate, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| run.click(translate, inputs=[audio_in, direction, mode], outputs=[out_text, timeline], api_name="translate") | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |