Update app.py

app.py

@@ -2,9 +2,7 @@ import gradio as gr
 import whisper
 import os
 from pydub import AudioSegment
-from transformers import AutoProcessor, AutoModelForCTC
-import torchaudio
-import torch
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
 
 # Mapping of model names to Whisper model sizes
 MODELS = {
@@ -12,12 +10,11 @@ MODELS = {
     "Base (Faster)": "base",
     "Small (Balanced)": "small",
     "Medium (Accurate)": "medium",
-    "Large (Most Accurate)": "large"
+    "Large (Most Accurate)": "large",
+    "Fine-Tuned Hindi": "yash-04/whisper-base-hindi",  # Hindi fine-tuned model
+    "Fine-Tuned Tamil": "mahimairaja/whisper-base-tamil"  # Tamil fine-tuned model
 }
 
-# Fine-tuned Sinhala model (using Hugging Face Transformers)
-SINHALA_MODEL = "IAmNotAnanth/wav2vec2-large-xls-r-300m-sinhala"
-
 # Mapping of full language names to language codes
 LANGUAGE_NAME_TO_CODE = {
     "Auto Detect": "Auto Detect",
@@ -122,78 +119,106 @@ LANGUAGE_NAME_TO_CODE = {
     "Sundanese": "su",
 }
 
-def transcribe_with_whisper(audio_file, language="Auto Detect", model_size="Base (Faster)"):
-    """Transcribe using OpenAI's Whisper models."""
-    model = whisper.load_model(MODELS[model_size])
+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 compatibility with Whisper
+    # 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")
     
-    # Transcribe the audio
-    if language == "Auto Detect":
-        result = model.transcribe(processed_audio_path, fp16=False)
-        detected_language = result.get("language", "unknown")
-    else:
-        language_code = LANGUAGE_NAME_TO_CODE.get(language, "en")
-        result = model.transcribe(processed_audio_path, language=language_code, fp16=False)
-        detected_language = language_code
-
+    # 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}\n\nTranscription:\n{result['text']}"
+    
+    return f"Detected Language: {detected_language}"
 
-def transcribe_with_sinhala_model(audio_file):
-    """Transcribe using the fine-tuned Sinhala Wav2Vec2 model."""
-    processor = AutoProcessor.from_pretrained(SINHALA_MODEL)
-    model = AutoModelForCTC.from_pretrained(SINHALA_MODEL)
+def transcribe_audio(audio_file, language="Auto Detect", model_size="Base (Faster)"):
+    """Transcribe the audio file."""
+    # Map language to fine-tuned model
+    language_to_model = {
+        "Hindi": "yash-04/whisper-base-hindi",
+        "Tamil": "mahimairaja/whisper-base-tamil",
+        # Add more mappings as needed
+    }
 
+    # Load the selected Whisper model
+    if language in language_to_model:
+        model_name = language_to_model[language]
+        model = WhisperForConditionalGeneration.from_pretrained(model_name)
+        processor = WhisperProcessor.from_pretrained(model_name)
+    else:
+        model = whisper.load_model(MODELS[model_size])
+        processor = None  # Use default Whisper processor
+
+    # 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")
 
-    audio_input, _ = torchaudio.load(processed_audio_path)
-    input_values = processor(audio_input.squeeze(), return_tensors="pt", sampling_rate=16000).input_values
-    logits = model(input_values).logits
-    predicted_ids = torch.argmax(logits, dim=-1)
+    # Transcribe the audio
+    if language == "Auto Detect":
+        if processor:
+            inputs = processor(processed_audio_path, return_tensors="pt", sampling_rate=16000)
+            result = model.generate(inputs.input_features)
+            transcription = processor.batch_decode(result, skip_special_tokens=True)[0]
+        else:
+            result = model.transcribe(processed_audio_path, fp16=False)
+            transcription = result["text"]
+        detected_language = result.get("language", "unknown")
+    else:
+        language_code = LANGUAGE_NAME_TO_CODE.get(language, "en")  # Default to English if not found
+        if processor:
+            inputs = processor(processed_audio_path, return_tensors="pt", sampling_rate=16000)
+            result = model.generate(inputs.input_features, language=language_code)
+            transcription = processor.batch_decode(result, skip_special_tokens=True)[0]
+        else:
+            result = model.transcribe(processed_audio_path, language=language_code, fp16=False)
+            transcription = result["text"]
+        detected_language = language_code
 
-    transcription = processor.batch_decode(predicted_ids)[0]
+    # Clean up processed audio file
     os.remove(processed_audio_path)
-    return f"Transcription:\n{transcription}"
 
-def transcribe_audio(audio_file, language="Auto Detect", model_size="Base (Faster)"):
-    """Wrapper to select the correct transcription method."""
-    if language == "Sinhala":
-        return transcribe_with_sinhala_model(audio_file)
-    else:
-        return transcribe_with_whisper(audio_file, language, model_size)
+    # Return transcription and detected language
+    return f"Detected Language: {detected_language}\n\nTranscription:\n{transcription}"
 
-# Gradio interface
+# Define the Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Audio Transcription and Language Detection")
     
-    transcribe_audio_input = gr.Audio(type="filepath", label="Upload Audio File")
-    language_dropdown = gr.Dropdown(
-        choices=list(LANGUAGE_NAME_TO_CODE.keys()),
-        label="Select Language",
-        value="Auto Detect"
-    )
-    model_dropdown = gr.Dropdown(
-        choices=list(MODELS.keys()),
-        label="Select Whisper Model",
-        value="Base (Faster)"
-    )
-    transcribe_output = gr.Textbox(label="Transcription")
-    transcribe_button = gr.Button("Transcribe Audio")
+    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")
     
-    def update_model_dropdown(language):
-        if language == "Sinhala":
-            return gr.update(interactive=False, value="Base (Faster)")  # Disable dropdown
-        return gr.update(interactive=True, value="Base (Faster)")
-
-    language_dropdown.change(update_model_dropdown, inputs=language_dropdown, outputs=model_dropdown)
+    with gr.Tab("Transcribe Audio"):
+        gr.Markdown("Upload an audio file, select a language (or choose 'Auto Detect'), and choose a model for transcription.")
+        transcribe_audio_input = gr.Audio(type="filepath", label="Upload Audio File")
+        language_dropdown = gr.Dropdown(
+            choices=list(LANGUAGE_NAME_TO_CODE.keys()),  # Full language names
+            label="Select Language",
+            value="Auto Detect"
+        )
+        model_dropdown = gr.Dropdown(
+            choices=list(MODELS.keys()),  # Model options
+            label="Select Model",
+            value="Base (Faster)"  # Default to "Base" model
+        )
+        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)
     transcribe_button.click(transcribe_audio, inputs=[transcribe_audio_input, language_dropdown, model_dropdown], outputs=transcribe_output)
 
-demo.launch()
+# Launch the Gradio interface
+demo.launch()