Update app.py

app.py

@@ -3,7 +3,7 @@ import whisper
 import torch
 import os
 from pydub import AudioSegment
-from transformers import AutoProcessor, AutoModelForCTC
+from transformers import pipeline
 
 # Mapping of model names to Whisper model sizes
 MODELS = {
@@ -14,13 +14,13 @@ MODELS = {
     "Large (Most Accurate)": "large"
 }
 
-# Fine-tuned Wav2Vec2 models for specific languages
-WAV2VEC2_MODELS = {
+# Fine-tuned models for specific languages
+FINE_TUNED_MODELS = {
     "Tamil": {
-        "processor": "Amrrs/wav2vec2-large-xlsr-53-tamil",
-        "model": "Amrrs/wav2vec2-large-xlsr-53-tamil"
+        "model": "vasista22/whisper-tamil-medium",
+        "language": "ta"
     },
-    # Add more Wav2Vec2 models for other languages here
+    # Add more fine-tuned models for other languages here
 }
 
 # Mapping of full language names to language codes
@@ -136,17 +136,21 @@ def transcribe_audio(audio_file, language="Auto Detect", model_size="Base (Faste
     audio.export(processed_audio_path, format="wav")
     
     # Load the appropriate model
-    if language in WAV2VEC2_MODELS:
-        # Use the fine-tuned Wav2Vec2 model for the selected language
-        processor = AutoProcessor.from_pretrained(WAV2VEC2_MODELS[language]["processor"])
-        model = AutoModelForCTC.from_pretrained(WAV2VEC2_MODELS[language]["model"])
-        
-        # Load audio and process
-        inputs = processor(AudioSegment.from_file(processed_audio_path).raw_data, sampling_rate=16000, return_tensors="pt")
-        with torch.no_grad():
-            logits = model(inputs.input_values).logits
-        predicted_ids = torch.argmax(logits, dim=-1)
-        transcription = processor.decode(predicted_ids[0])
+    if language in FINE_TUNED_MODELS:
+        # Use the fine-tuned Whisper model for the selected language
+        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        transcribe = pipeline(
+            task="automatic-speech-recognition",
+            model=FINE_TUNED_MODELS[language]["model"],
+            chunk_length_s=30,
+            device=device
+        )
+        transcribe.model.config.forced_decoder_ids = transcribe.tokenizer.get_decoder_prompt_ids(
+            language=FINE_TUNED_MODELS[language]["language"],
+            task="transcribe"
+        )
+        result = transcribe(processed_audio_path)
+        transcription = result["text"]
         detected_language = language
     else:
         # Use the selected Whisper model
@@ -192,7 +196,7 @@ with gr.Blocks() as demo:
     
     # Update model dropdown based on language selection
     def update_model_dropdown(language):
-        if language in WAV2VEC2_MODELS:
+        if language in FINE_TUNED_MODELS:
             return gr.Dropdown(interactive=False, value=f"Fine-Tuned {language} Model")
         else:
             return gr.Dropdown(choices=list(MODELS.keys()), interactive=True, value="Base (Faster)")