Update app.py

app.py

@@ -1,7 +1,12 @@
 import gradio as gr
 import whisper
+import torch
 import os
 from pydub import AudioSegment
+from transformers import pipeline
+
+# Ensure compatible versions of torch and transformers are installed
+# Run: pip install torch==1.13.1 transformers==4.26.1
 
 # Mapping of model names to Whisper model sizes
 MODELS = {
@@ -12,6 +17,15 @@ MODELS = {
     "Large (Most Accurate)": "large"
 }
 
+# Fine-tuned models for specific languages
+FINE_TUNED_MODELS = {
+    "Tamil": {
+        "model": "vasista22/whisper-tamil-medium",
+        "language": "ta"
+    },
+    # Add more fine-tuned models for other languages here
+}
+
 # Mapping of full language names to language codes
 LANGUAGE_NAME_TO_CODE = {
     "Auto Detect": "Auto Detect",
@@ -116,61 +130,55 @@ LANGUAGE_NAME_TO_CODE = {
     "Sundanese": "su",
 }
 
-def detect_language(audio_file):
-    """Detect the language of the audio file."""
-    # Load the Whisper model (use "base" for faster detection)
-    model = whisper.load_model("base")
-    
-    # Convert audio to 16kHz mono for better compatibility with Whisper
-    audio = AudioSegment.from_file(audio_file)
-    audio = audio.set_frame_rate(16000).set_channels(1)
-    processed_audio_path = "processed_audio.wav"
-    audio.export(processed_audio_path, format="wav")
-    
-    # Detect the language
-    result = model.transcribe(processed_audio_path, task="detect_language", fp16=False)
-    detected_language = result.get("language", "unknown")
-    
-    # Clean up processed audio file
-    os.remove(processed_audio_path)
-    
-    return f"Detected Language: {detected_language}"
-
 def transcribe_audio(audio_file, language="Auto Detect", model_size="Base (Faster)"):
     """Transcribe the audio file."""
-    # Load the selected Whisper model
-    model = whisper.load_model(MODELS[model_size])
-    
-    # Convert audio to 16kHz mono for better compatibility with Whisper
+    # Convert audio to 16kHz mono for better compatibility
     audio = AudioSegment.from_file(audio_file)
     audio = audio.set_frame_rate(16000).set_channels(1)
     processed_audio_path = "processed_audio.wav"
     audio.export(processed_audio_path, format="wav")
     
-    # Transcribe the audio
-    if language == "Auto Detect":
-        result = model.transcribe(processed_audio_path, fp16=False)  # Auto-detect language
-        detected_language = result.get("language", "unknown")
+    # Load the appropriate model
+    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:
-        language_code = LANGUAGE_NAME_TO_CODE.get(language, "en")  # Default to English if not found
-        result = model.transcribe(processed_audio_path, language=language_code, fp16=False)
-        detected_language = language_code
+        # Use the selected Whisper model
+        model = whisper.load_model(MODELS[model_size])
+        
+        # Transcribe the audio
+        if language == "Auto Detect":
+            result = model.transcribe(processed_audio_path, fp16=False)  # Auto-detect language
+            detected_language = result.get("language", "unknown")
+        else:
+            language_code = LANGUAGE_NAME_TO_CODE.get(language, "en")  # Default to English if not found
+            result = model.transcribe(processed_audio_path, language=language_code, fp16=False)
+            detected_language = language_code
+        
+        transcription = result["text"]
     
     # Clean up processed audio file
     os.remove(processed_audio_path)
     
     # Return transcription and detected language
-    return f"Detected Language: {detected_language}\n\nTranscription:\n{result['text']}"
+    return f"Detected Language: {detected_language}\n\nTranscription:\n{transcription}"
 
 # Define the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Audio Transcription and Language Detection")
-    
-    with gr.Tab("Detect Language"):
-        gr.Markdown("Upload an audio file to detect its language.")
-        detect_audio_input = gr.Audio(type="filepath", label="Upload Audio File")
-        detect_language_output = gr.Textbox(label="Detected Language")
-        detect_button = gr.Button("Detect Language")
+    gr.Markdown("# Audio Transcription with Fine-Tuned Models")
     
     with gr.Tab("Transcribe Audio"):
         gr.Markdown("Upload an audio file, select a language (or choose 'Auto Detect'), and choose a model for transcription.")
@@ -183,13 +191,24 @@ with gr.Blocks() as demo:
         model_dropdown = gr.Dropdown(
             choices=list(MODELS.keys()),  # Model options
             label="Select Model",
-            value="Base (Faster)"  # Default to "Base" model
+            value="Base (Faster)",  # Default to "Base" model
+            interactive=True  # Allow model selection by default
         )
         transcribe_output = gr.Textbox(label="Transcription and Detected Language")
         transcribe_button = gr.Button("Transcribe Audio")
     
-    # Link buttons to functions
-    detect_button.click(detect_language, inputs=detect_audio_input, outputs=detect_language_output)
+    # Update model dropdown based on language selection
+    def update_model_dropdown(language):
+        if language in FINE_TUNED_MODELS:
+            # Add "Fine-Tuned Model" to the dropdown choices and disable it
+            return gr.Dropdown(choices=["Fine-Tuned Model"], value="Fine-Tuned Model", interactive=False)
+        else:
+            # Reset the dropdown to standard Whisper models
+            return gr.Dropdown(choices=list(MODELS.keys()), value="Base (Faster)", interactive=True)
+    
+    language_dropdown.change(update_model_dropdown, inputs=language_dropdown, outputs=model_dropdown)
+    
+    # Link button to function
     transcribe_button.click(transcribe_audio, inputs=[transcribe_audio_input, language_dropdown, model_dropdown], outputs=transcribe_output)
 
 # Launch the Gradio interface