add punctuations

app.py

@@ -26,7 +26,7 @@ from pathlib import Path
 import gradio as gr
 
 from decode import decode
-from model import get_pretrained_model, get_vad, language_to_models
+from model import get_pretrained_model, get_vad, language_to_models, get_punct_model
 
 title = "# Next-gen Kaldi: Generate subtitles for videos"
 
@@ -89,6 +89,7 @@ def show_file_info(in_filename: str):
 def process_uploaded_video_file(
     language: str,
     repo_id: str,
+    add_punctuation: str,
     in_filename: str,
 ):
     if in_filename is None or in_filename == "":
@@ -105,13 +106,14 @@ def process_uploaded_video_file(
 
     logging.info(f"Processing uploaded file: {in_filename}")
 
-    ans = process(language, repo_id, in_filename)
+    ans = process(language, repo_id, add_punctuation, in_filename)
     return (in_filename, ans[0]), ans[0], ans[1], ans[2]
 
 
 def process_uploaded_audio_file(
     language: str,
     repo_id: str,
+    add_punctuation: str,
     in_filename: str,
 ):
     if in_filename is None or in_filename == "":
@@ -131,11 +133,15 @@ def process_uploaded_audio_file(
     return process(language, repo_id, in_filename)
 
 
-def process(language: str, repo_id: str, in_filename: str):
+def process(language: str, repo_id: str, add_punctuation: str, in_filename: str):
     recognizer = get_pretrained_model(repo_id)
     vad = get_vad()
+    if add_punctuation == "Yes":
+        punct = get_punct_model()
+    else:
+        punct = None
 
-    result = decode(recognizer, vad, in_filename)
+    result = decode(recognizer, vad, punct, in_filename)
     logging.info(result)
 
     srt_filename = Path(in_filename).with_suffix(".srt")
@@ -176,6 +182,11 @@ with demo:
         inputs=language_radio,
         outputs=model_dropdown,
     )
+    punct_radio = gr.Radio(
+        label="Whether to add punctuation",
+        choices=["Yes", "No"],
+        value="Yes",
+    )
 
     with gr.Tabs():
         with gr.TabItem("Upload video from disk"):
@@ -218,6 +229,7 @@ with demo:
             inputs=[
                 language_radio,
                 model_dropdown,
+                punct_radio,
                 uploaded_video_file,
             ],
             outputs=[
@@ -233,6 +245,7 @@ with demo:
             inputs=[
                 language_radio,
                 model_dropdown,
+                punct_radio,
                 uploaded_audio_file,
             ],
             outputs=[

decode.py

@@ -48,6 +48,7 @@ class Segment:
 def decode(
     recognizer: sherpa_onnx.OfflineRecognizer,
     vad: sherpa_onnx.VoiceActivityDetector,
+    punct: Optional[sherpa_onnx.OfflinePunctuation],
     filename: str,
 ) -> str:
     ffmpeg_cmd = [
@@ -114,6 +115,8 @@ def decode(
 
         for seg, stream in zip(segments, streams):
             seg.text = stream.result.text.strip()
+            if punct is not None:
+                seg.text = punct.add_punctuation(seg.text)
             segment_list.append(seg)
 
     return "\n\n".join(f"{i}\n{seg}" for i, seg in enumerate(segment_list, 1))

model.py

@@ -168,6 +168,21 @@ def _get_russian_pre_trained_model(repo_id: str) -> sherpa_onnx.OfflineRecognize
     return recognizer
 
 
+@lru_cache(maxsize=2)
+def get_punct_model() -> sherpa_onnx.OfflinePunctuation:
+    model = _get_nn_model_filename(
+        repo_id="csukuangfj/sherpa-onnx-punct-ct-transformer-zh-en-vocab272727-2024-04-12",
+        filename="model.onnx",
+        subfolder=".",
+    )
+    config = sherpa_onnx.OfflinePunctuationConfig(
+        model=sherpa_onnx.OfflinePunctuationModelConfig(ct_transformer=model),
+    )
+
+    punct = sherpa_onnx.OfflinePunctuation(config)
+    return punct
+
+
 @lru_cache(maxsize=2)
 def get_vad() -> sherpa_onnx.VoiceActivityDetector:
     vad_model = _get_nn_model_filename(