add llm asr

app.py

@@ -9,12 +9,17 @@ from fireredasr.models.fireredasr import FireRedAsr
 
 
 asr_model_aed = None
+asr_model_llm = None
 
 
-def init_model(model_dir_aed):
+def init_model(model_dir_aed, model_dir_llm):
     global asr_model_aed
+    global asr_model_llm
     if asr_model_aed is None:
         asr_model_aed = FireRedAsr.from_pretrained("aed", model_dir_aed)
+    if asr_model_llm is None:
+        asr_model_llm = FireRedAsr.from_pretrained("llm", model_dir_llm)
+
 
 @spaces.GPU(duration=20)
 def asr_inference(audio_file):
@@ -43,6 +48,30 @@ def asr_inference(audio_file):
     return text_output
 
 
+@spaces.GPU(duration=30)
+def asr_inference_llm(audio_file):
+    if not audio_file:
+        return "Please upload a wav file"
+    batch_uttid = ["demo"]
+    batch_wav_path = [audio_file]
+    results = asr_model_llm.transcribe(
+        batch_uttid,
+        batch_wav_path,
+        {
+            "use_gpu": True,
+            "beam_size": 3,
+            "nbest": 1,
+            "decode_max_len": 0,
+            "decode_min_len": 0,
+            "repetition_penalty": 3.0,
+            "llm_length_penalty": 1.0,
+            "temperature": 1.0
+        }
+    )
+    text_output = results[0]["text"]
+    return text_output
+
+
 with gr.Blocks(title="FireRedASR") as demo:
     gr.HTML(
         "<h1 style='text-align: center'>FireRedASR Demo</h1>"
@@ -53,10 +82,12 @@ with gr.Blocks(title="FireRedASR") as demo:
         with gr.Column():
             #audio_file = gr.Audio(label="Upload Audio", sources=["upload", "microphone"], type="filepath")
             audio_file = gr.Audio(label="Upload wav file", sources=["upload"], type="filepath")
-            asr_button = gr.Button("Start Recognition", variant="primary")
 
         with gr.Column():
-            text_output = gr.Textbox(label="Model Result", interactive=False, lines=6, max_lines=12)
+            asr_button = gr.Button("Start Recognition (FireRedASR-AED-L)", variant="primary")
+            text_output = gr.Textbox(label="Model Result (FireRedASR-AED-L)", interactive=False, lines=3, max_lines=12)
+            asr_button_llm = gr.Button("Start Recognition (FireRedASR-LLM-L)", variant="primary")
+            text_output_llm = gr.Textbox(label="Model Result (FireRedASR-LLM-L)", interactive=False, lines=3, max_lines=12)
 
     asr_button.click(
         fn=asr_inference,
@@ -64,13 +95,21 @@ with gr.Blocks(title="FireRedASR") as demo:
         outputs=[text_output]
     )
 
+    asr_button_llm.click(
+        fn=asr_inference_llm,
+        inputs=[audio_file],
+        outputs=[text_output_llm]
+    )
+
 
 if __name__ == "__main__":
     # Download model
     local_dir='pretrained_models/FireRedASR-AED-L'
     snapshot_download(repo_id='FireRedTeam/FireRedASR-AED-L', local_dir=local_dir)
+    local_dir_llm='pretrained_models/FireRedASR-LLM-L'
+    snapshot_download(repo_id='FireRedTeam/FireRedASR-LLM-L', local_dir=local_dir_llm)
     # Init model
-    init_model(local_dir)
+    init_model(local_dir, local_dir_llm)
     # UI
     demo.queue()
     demo.launch()