FireRedTeam/FireRedASR

1import sys2 3import gradio as gr4import spaces5from huggingface_hub import snapshot_download6 7sys.path.append("./fireredasr")8from fireredasr.models.fireredasr import FireRedAsr9 10 11asr_model_aed = None12asr_model_llm = None13 14 15def init_model(model_dir_aed, model_dir_llm):16    global asr_model_aed17    global asr_model_llm18    if asr_model_aed is None:19        asr_model_aed = FireRedAsr.from_pretrained("aed", model_dir_aed)20    if asr_model_llm is None:21        asr_model_llm = FireRedAsr.from_pretrained("llm", model_dir_llm)22 23 24@spaces.GPU(duration=20)25def asr_inference(audio_file):26    if not audio_file:27        return "Please upload a wav file"28    batch_uttid = ["demo"]29    batch_wav_path = [audio_file]30    results = asr_model_aed.transcribe(31        batch_uttid,32        batch_wav_path,33        {34            "use_gpu": True,35            "beam_size": 3,36            "nbest": 1,37            "decode_max_len": 0,38            "softmax_smoothing": 1.25,39            "aed_length_penalty": 0.6,40            "eos_penalty": 1.0,41            #"decode_min_len": args.decode_min_len,42            #"repetition_penalty": args.repetition_penalty,43            #"llm_length_penalty": args.llm_length_penalty,44            #"temperature": args.temperature45        }46    )47    text_output = results[0]["text"]48    return text_output49 50 51@spaces.GPU(duration=30)52def asr_inference_llm(audio_file):53    if not audio_file:54        return "Please upload a wav file"55    batch_uttid = ["demo"]56    batch_wav_path = [audio_file]57    results = asr_model_llm.transcribe(58        batch_uttid,59        batch_wav_path,60        {61            "use_gpu": True,62            "beam_size": 3,63            "nbest": 1,64            "decode_max_len": 0,65            "decode_min_len": 0,66            "repetition_penalty": 3.0,67            "llm_length_penalty": 1.0,68            "temperature": 1.069        }70    )71    text_output = results[0]["text"]72    return text_output73 74 75with gr.Blocks(title="FireRedASR") as demo:76    gr.HTML(77        "<h1 style='text-align: center'>FireRedASR Demo</h1>"78    )79    gr.Markdown("Upload an audio file (wav) to get speech-to-text results.")80 81    with gr.Row():82        with gr.Column():83            #audio_file = gr.Audio(label="Upload Audio", sources=["upload", "microphone"], type="filepath")84            audio_file = gr.Audio(label="Upload wav file", sources=["upload"], type="filepath")85 86        with gr.Column():87            asr_button = gr.Button("Start Recognition (FireRedASR-AED-L)", variant="primary")88            text_output = gr.Textbox(label="Model Result (FireRedASR-AED-L)", interactive=False, lines=3, max_lines=12)89            asr_button_llm = gr.Button("Start Recognition (FireRedASR-LLM-L)", variant="primary")90            text_output_llm = gr.Textbox(label="Model Result (FireRedASR-LLM-L)", interactive=False, lines=3, max_lines=12)91 92    asr_button.click(93        fn=asr_inference,94        inputs=[audio_file],95        outputs=[text_output]96    )97 98    asr_button_llm.click(99        fn=asr_inference_llm,100        inputs=[audio_file],101        outputs=[text_output_llm]102    )103 104 105if __name__ == "__main__":106    # Download model107    local_dir='pretrained_models/FireRedASR-AED-L'108    snapshot_download(repo_id='FireRedTeam/FireRedASR-AED-L', local_dir=local_dir)109    local_dir_llm='pretrained_models/FireRedASR-LLM-L'110    snapshot_download(repo_id='FireRedTeam/FireRedASR-LLM-L', local_dir=local_dir_llm)111    local_dir_qwen='pretrained_models/FireRedASR-LLM-L/Qwen2-7B-Instruct'112    snapshot_download(repo_id='Qwen/Qwen2-7B-Instruct', local_dir=local_dir_qwen)113    # Init model114    init_model(local_dir, local_dir_llm)115    # UI116    demo.queue()117    demo.launch()118