import gradio as gr
import torch
import torchaudio

# 定义模型路径
model_path = "https://huggingface.co/Tele-AI/TeleSpeech-ASR1.0/blob/main/large.pt"

# 下载模型
model = torch.hub.load('snakers4/silero-models', 'silero_stt', language='en', device='cpu')
torch.hub.download_url_to_file(model_path, 'large.pt')

# 加载模型
model.load_state_dict(torch.load('large.pt', map_location=torch.device('cpu')))
model.eval()

# 定义处理函数
def transcribe(audio):
    waveform, sample_rate = torchaudio.load(audio.name)
    waveform = waveform.unsqueeze(0)
    with torch.no_grad():
        output = model(waveform)
    transcription = output[0]
    return transcription

# 创建 Gradio 界面
iface = gr.Interface(
    fn=transcribe, 
    inputs=gr.inputs.Audio(source="microphone", type="filepath"), 
    outputs="text",
    title="TeleSpeech ASR",
    description="Upload an audio file or record your voice to transcribe speech to text using the TeleSpeech ASR model."
)

iface.launch()