feat: 增加双语字幕

decode.py

@@ -19,6 +19,7 @@ import subprocess
 from dataclasses import dataclass
 from datetime import timedelta
 from typing import Optional
+from transformers import pipeline, MarianMTModel, MarianTokenizer
 
 import numpy as np
 import sherpa_onnx
@@ -122,7 +123,9 @@ def decode(
             recognizer.decode_stream(s)
 
         for seg, stream in zip(segments, streams):
-            seg.text = stream.result.text.strip()
+            en_text = stream.result.text.strip()
+            cn_text = _llm_translator.translate(en_text)
+            seg.text = en_text +"\n"+cn_text
             if len(seg.text) == 0:
                 logging.info("Skip empty segment")
                 continue
@@ -143,3 +146,32 @@ def decode(
         all_text = punct.add_punctuation(all_text)
 
     return "\n\n".join(f"{i}\n{seg}" for i, seg in enumerate(segment_list, 1)), all_text
+
+
+
+
+def translate_en_to_cn(src_text: str, ) -> str:
+
+    model_name = "Helsinki-NLP/opus-mt-en-zh"
+    tokenizer = MarianTokenizer.from_pretrained(model_name)
+    model = MarianMTModel.from_pretrained(model_name)
+    translated = model.generate(**tokenizer(src_text, return_tensors="pt", padding=True))
+    res = [tokenizer.decode(t, skip_special_tokens=True) for t in translated]
+    return res
+
+
+class LLMTranslator:
+    _tokenizer: MarianTokenizer
+    _model: MarianMTModel
+    def __init__(self):
+        model_name = "Helsinki-NLP/opus-mt-en-zh"
+        self._tokenizer = MarianTokenizer.from_pretrained(model_name)
+        self._model = MarianMTModel.from_pretrained(model_name)
+
+    def translate(self, src_text: str) -> str:
+        translated = self._model.generate(**self._tokenizer(src_text, return_tensors="pt", padding=True))
+        res = [self._tokenizer.decode(t, skip_special_tokens=True) for t in translated]
+        return res
+
+
+_llm_translator = LLMTranslator()