Update app.py

app.py

@@ -1,45 +1,140 @@
-import gradio as gr
-import torch
-import librosa
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
-
-# Load Wav2Vec2 Model
-MODEL_NAME = "facebook/wav2vec2-large-960h"
-processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
-model = Wav2Vec2ForCTC.from_pretrained(MODEL_NAME)
-
-def transcribe(audio_file):
-    """
-    Transcribes speech from an uploaded audio file or live microphone input.
-    """
-    try:
-        # Load and convert audio to 16kHz
-        audio, rate = librosa.load(audio_file, sr=16000)
-
-        # Convert audio to tensor format for Wav2Vec
-        input_values = processor(audio, sampling_rate=16000, return_tensors="pt").input_values
-
-        # Run the model for transcription
-        with torch.no_grad():
-            logits = model(input_values).logits
-
-        # Convert predicted tokens into text
-        predicted_ids = torch.argmax(logits, dim=-1)
-        transcription = processor.batch_decode(predicted_ids)[0]
-
-        return transcription
-
-    except Exception as e:
-        return "Error processing file"
-
-# UI Build
-interface = gr.Interface(
-    fn=transcribe,  
-    inputs=gr.Audio(sources=["microphone", "upload"], type="filepath", label="Speak or Upload Audio"),
-    outputs="text",
-    title="Wav2Vec2 Speech-to-Text Transcription",
-    description="Speak into your microphone or upload an audio file to get an automatic transcription.",
-    live=True  # Real-time microphone processing
-)
-
-interface.launch(share=True)
+import gradio as gr
+import torch
+import librosa
+from transformers import Wav2Vec2Processor, AutoModelForCTC
+import zipfile
+import os
+import firebase_admin
+from firebase_admin import credentials, firestore
+from datetime import datetime
+
+# 🔹 Initialize Firebase
+cred = credentials.Certificate('firebase_credentials.json')  # Your Firebase JSON key file
+firebase_admin.initialize_app(cred)
+db = firestore.client()
+
+# Load the ASR model and processor
+MODEL_NAME = "eleferrand/xlsr53_Amis"
+processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
+model = AutoModelForCTC.from_pretrained(MODEL_NAME)
+
+def transcribe(audio_file):
+    """
+    Transcribes the audio file using the loaded ASR model.
+    Returns the transcription string.
+    """
+    try:
+        # Load and resample the audio to 16kHz
+        audio, rate = librosa.load(audio_file, sr=16000)
+        # Prepare the input tensor for the model
+        input_values = processor(audio, sampling_rate=16000, return_tensors="pt").input_values
+
+        # Get model predictions (logits) and decode to text
+        with torch.no_grad():
+            logits = model(input_values).logits
+        predicted_ids = torch.argmax(logits, dim=-1)
+        transcription = processor.batch_decode(predicted_ids)[0]
+        return transcription.replace("[UNK]", "")
+    except Exception as e:
+        return f"Error processing file: {e}"
+
+def transcribe_both(audio_file):
+    """
+    Calls the transcribe function and returns the transcription
+    for both the original (read-only) and the corrected (editable) textboxes.
+    """
+    transcription = transcribe(audio_file)
+    return transcription, transcription
+
+def store_correction(original_transcription, corrected_transcription):
+    """
+    Stores the original and corrected transcription in Firestore.
+    """
+    try:
+        correction_data = {
+            'original_text': original_transcription,
+            'corrected_text': corrected_transcription,
+            'timestamp': datetime.now().isoformat()
+        }
+        db.collection('transcription_corrections').add(correction_data)
+        return "✅ Correction saved successfully!"
+    except Exception as e:
+        return f"⚠️ Error saving correction: {e}"
+
+def prepare_download(audio_file, original_transcription, corrected_transcription):
+    """
+    Prepares a ZIP file containing:
+      - The uploaded audio file (saved as audio.wav)
+      - A text file with the original transcription
+      - A text file with the corrected transcription
+    Returns the path to the ZIP file.
+    """
+    if audio_file is None:
+        return None
+
+    zip_filename = "results.zip"
+    with zipfile.ZipFile(zip_filename, "w") as zf:
+        # Add the audio file (saved as audio.wav in the zip)
+        if os.path.exists(audio_file):
+            zf.write(audio_file, arcname="audio.wav")
+        else:
+            print("Audio file not found:", audio_file)
+
+        # Create and add the original transcription file
+        orig_txt = "original_transcription.txt"
+        with open(orig_txt, "w", encoding="utf-8") as f:
+            f.write(original_transcription)
+        zf.write(orig_txt, arcname="original_transcription.txt")
+        os.remove(orig_txt)
+
+        # Create and add the corrected transcription file
+        corr_txt = "corrected_transcription.txt"
+        with open(corr_txt, "w", encoding="utf-8") as f:
+            f.write(corrected_transcription)
+        zf.write(corr_txt, arcname="corrected_transcription.txt")
+        os.remove(corr_txt)
+    return zip_filename
+
+# Build the Gradio Blocks interface
+with gr.Blocks() as demo:
+    gr.Markdown("# ASR Demo with Editable Transcription, Firestore Storage, and Download")
+    
+    with gr.Row():
+        audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Upload or Record Audio")
+        transcribe_button = gr.Button("Transcribe Audio")
+    
+    with gr.Row():
+        # The original transcription is displayed (non-editable)
+        original_text = gr.Textbox(label="Original Transcription", interactive=False)
+        # The corrected transcription is pre-filled with the original, but remains editable.
+        corrected_text = gr.Textbox(label="Corrected Transcription", interactive=True)
+
+    save_button = gr.Button("Save Correction to Database")
+    save_status = gr.Textbox(label="Save Status", interactive=False)
+
+    download_button = gr.Button("Download Results (ZIP)")
+    download_output = gr.File(label="Download ZIP")
+    
+    # When the transcribe button is clicked, update both textboxes with the transcription.
+    transcribe_button.click(
+        fn=transcribe_both, 
+        inputs=audio_input, 
+        outputs=[original_text, corrected_text]
+    )
+    
+    # When the "Save Correction" button is clicked, store the corrected transcription in Firestore.
+    save_button.click(
+        fn=store_correction, 
+        inputs=[original_text, corrected_text], 
+        outputs=save_status
+    )
+
+    # When the download button is clicked, package the audio file and both transcriptions into a zip.
+    download_button.click(
+        fn=prepare_download, 
+        inputs=[audio_input, original_text, corrected_text], 
+        outputs=download_output
+    )
+
+# Launch the demo
+demo.launch(share=True)