Update app.py

app.py

@@ -1,50 +1,61 @@
-import gradio as gr
 import torch
-from transformers import WhisperProcessor, WhisperForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
+import torchaudio
+import gradio as gr
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load Whisper fine-tuned Odia model
+whisper_model_path = "./whisper-odia-final"  # Change if needed
+processor = WhisperProcessor.from_pretrained(whisper_model_path)
+model = WhisperForConditionalGeneration.from_pretrained(whisper_model_path).to(device)
+
+# Load IndicTrans2 multilingual model
+trans_model_id = "ai4bharat/indictrans2-en-indic-dist-200M "
+translator_tokenizer = AutoTokenizer.from_pretrained(trans_model_id, use_fast=False)
+translator_model = AutoModelForSeq2SeqLM.from_pretrained(trans_model_id).to(device)
+
+# Translation function with language tags
+def translate_to_english(text):
+    if not text.strip():
+        return ""
 
-# Load Whisper Odia ASR model (fine-tuned)
-asr_model = WhisperForConditionalGeneration.from_pretrained("ashutoshpattnaik50/whisper-small-odia-finetuned")
-processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+    # Add source and target language tokens
+    text_with_lang = f"<2en> {text.strip()}"
+    inputs = translator_tokenizer(text_with_lang, return_tensors="pt", padding=True).to(device)
 
-# Load IndicTrans2 Odia-to-English model
-translator = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-en")
-translator_tokenizer = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-en")
+    output = translator_model.generate(**inputs, max_length=256)
+    translated = translator_tokenizer.batch_decode(output, skip_special_tokens=True)[0]
+    return translated
 
-# Transcription + Translation function
-def transcribe_and_translate(audio_path):
-    # Load and preprocess audio
-    from datasets import load_dataset, Audio
-    import numpy as np
-    import torchaudio
+# ASR + Translation Pipeline
+def transcribe(audio_path):
+    if audio_path is None:
+        return "No audio received.", ""
 
-    speech_array, sampling_rate = torchaudio.load(audio_path)
-    if sampling_rate != 16000:
-        resampler = torchaudio.transforms.Resample(sampling_rate, 16000)
-        speech_array = resampler(speech_array)
-    
-    input_features = processor(speech_array.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features
+    speech, sr = torchaudio.load(audio_path)
+    if sr != 16000:
+        speech = torchaudio.functional.resample(speech, sr, 16000)
 
-    predicted_ids = asr_model.generate(input_features)
-    transcription = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+    input_features = processor(speech.squeeze().numpy(), sampling_rate=16000, return_tensors="pt").input_features.to(device)
 
-    # Translate to English
-    inputs = translator_tokenizer(transcription, return_tensors="pt")
-    output_tokens = translator.generate(**inputs)
-    translated = translator_tokenizer.batch_decode(output_tokens, skip_special_tokens=True)[0]
+    predicted_ids = model.generate(input_features)
+    odia_text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
 
-    return transcription.strip(), translated.strip()
+    english_text = translate_to_english(odia_text)
+    return odia_text, english_text
 
-# Gradio Interface
+# Gradio UI
 interface = gr.Interface(
-    fn=transcribe_and_translate,
-    inputs=gr.Audio(type="filepath", label="🎤 Record or Upload Odia Audio"),
+    fn=transcribe,
+    inputs=gr.Audio(source="microphone", type="filepath", label="🎤 Record or Upload Odia Audio"),
     outputs=[
         gr.Textbox(label="📝 Odia Transcription"),
         gr.Textbox(label="🌐 English Translation")
     ],
-    title="Odia Whisper ASR + Translator",
-    description="Speak in Odia and get instant transcription + English translation using Whisper and IndicTrans2.",
+    title="Whisper Odia ASR + Translation",
+    description="🎙️ Speak in Odia → Get Odia transcription → Get English translation using IndicTrans2"
 )
 
-if __name__ == "__main__":
-    interface.launch()
+interface.launch()