Spaces:
Running
Running
| from __future__ import annotations | |
| import torch | |
| import torchaudio | |
| import gradio as gr | |
| import spaces | |
| from transformers import AutoModel | |
| DESCRIPTION = "IndicConformer-600M Multilingual ASR (CTC + RNNT)" | |
| LANGUAGE_NAME_TO_CODE = { | |
| "Assamese": "as", "Bengali": "bn", "Bodo": "br", "Dogri": "doi", | |
| "Gujarati": "gu", "Hindi": "hi", "Kannada": "kn", "Kashmiri": "ks", | |
| "Konkani": "kok", "Maithili": "mai", "Malayalam": "ml", "Manipuri": "mni", | |
| "Marathi": "mr", "Nepali": "ne", "Odia": "or", "Punjabi": "pa", | |
| "Sanskrit": "sa", "Santali": "sat", "Sindhi": "sd", "Tamil": "ta", | |
| "Telugu": "te", "Urdu": "ur" | |
| } | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # Load Indic Conformer model (assumes custom forward handles decoding strategy) | |
| model = AutoModel.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True).to(device) | |
| model.eval() | |
| def transcribe_ctc_and_rnnt(audio_path, language_name): | |
| lang_code = LANGUAGE_NAME_TO_CODE[language_name] | |
| # Load and preprocess audio | |
| waveform, sr = torchaudio.load(audio_path) | |
| waveform = waveform.mean(dim=0, keepdim=True) if waveform.shape[0] > 1 else waveform | |
| waveform = torchaudio.functional.resample(waveform, sr, 16000).to(device) | |
| try: | |
| # Assume model's forward method takes waveform, language code, and decoding type | |
| with torch.no_grad(): | |
| transcription_ctc = model(waveform, lang_code, "ctc") | |
| transcription_rnnt = model(waveform, lang_code, "rnnt") | |
| except Exception as e: | |
| return f"Error: {str(e)}", "" | |
| return transcription_ctc.strip(), transcription_rnnt.strip() | |
| # Gradio UI | |
| with gr.Blocks() as demo: | |
| gr.Markdown(f"## {DESCRIPTION}") | |
| with gr.Row(): | |
| with gr.Column(): | |
| audio = gr.Audio(label="Upload or Record Audio", type="filepath") | |
| lang = gr.Dropdown( | |
| label="Select Language", | |
| choices=list(LANGUAGE_NAME_TO_CODE.keys()), | |
| value="Hindi" | |
| ) | |
| transcribe_btn = gr.Button("Transcribe (CTC + RNNT)") | |
| with gr.Column(): | |
| gr.Markdown("### CTC Transcription") | |
| ctc_output = gr.Textbox(lines=3) | |
| gr.Markdown("### RNNT Transcription") | |
| rnnt_output = gr.Textbox(lines=3) | |
| transcribe_btn.click(fn=transcribe_ctc_and_rnnt, inputs=[audio, lang], outputs=[ctc_output, rnnt_output], api_name="transcribe") | |
| if __name__ == "__main__": | |
| demo.queue().launch() | |