Update app.py (#2)

- Update app.py (ecaaa4ba11e5d53ea5ca74d2738df3fb1058d49a)


Co-authored-by: Hunter S <[email protected]>

app.py

@@ -7,7 +7,7 @@ import os
 import firebase_admin
 from firebase_admin import credentials, firestore
 from datetime import datetime
-import json 
+import json
 
 # Initialize Firebase
 firebase_config = json.loads(os.environ.get('firebase_creds'))
@@ -15,7 +15,6 @@ cred = credentials.Certificate(firebase_config)  # 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)
@@ -43,21 +42,46 @@ def transcribe(audio_file):
 
 def transcribe_both(audio_file):
     """
-    Calls the transcribe function and returns the transcription
-    for both the original (read-only) and the corrected (editable) textboxes.
+    Transcribes the audio and returns:
+      - the original transcription (for the non-editable textbox),
+      - the transcription (pre-filled for the editable textbox), and
+      - the processing time (in seconds).
     """
+    start_time = datetime.now()
     transcription = transcribe(audio_file)
-    return transcription, transcription
+    processing_time = (datetime.now() - start_time).total_seconds()
+    return transcription, transcription, processing_time
 
-def store_correction(original_transcription, corrected_transcription):
+def store_correction(original_transcription, corrected_transcription, audio_file, processing_time):
     """
-    Stores the original and corrected transcription in Firestore.
+    Stores the transcriptions and additional metadata in Firestore.
+    Saves:
+      - original & corrected text,
+      - timestamp,
+      - processing time,
+      - audio metadata (duration & file size, if available),
+      - a placeholder for the audio URL, and
+      - the model name.
     """
     try:
+        audio_metadata = {}
+        if audio_file and os.path.exists(audio_file):
+            # Load audio for metadata calculations
+            audio, sr = librosa.load(audio_file, sr=16000)
+            duration = librosa.get_duration(y=audio, sr=sr)
+            file_size = os.path.getsize(audio_file)
+            audio_metadata = {
+                'duration': duration,
+                'file_size': file_size
+            }
         correction_data = {
             'original_text': original_transcription,
             'corrected_text': corrected_transcription,
-            'timestamp': datetime.now().isoformat()
+            'timestamp': datetime.now().isoformat(),
+            'processing_time': processing_time,
+            'audio_metadata': audio_metadata,
+            'audio_url': None,
+            'model_name': MODEL_NAME
         }
         db.collection('transcription_corrections').add(correction_data)
         return "Correction saved successfully!"
@@ -67,30 +91,30 @@ def store_correction(original_transcription, corrected_transcription):
 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.
+      - The uploaded audio file (as audio.wav),
+      - a text file with the original transcription, and
+      - a text file with the corrected transcription.
+    Returns the ZIP file's path.
     """
     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)
+        # Add the audio file (renamed inside 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
+        # Add the original transcription as a text 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
+        # Add the corrected transcription as a text file
         corr_txt = "corrected_transcription.txt"
         with open(corr_txt, "w", encoding="utf-8") as f:
             f.write(corrected_transcription)
@@ -98,46 +122,47 @@ def prepare_download(audio_file, original_transcription, corrected_transcription
         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
-    )
+# Build the Gradio Blocks interface with improved styling
+with gr.Blocks(css="""
+    .container { max-width: 800px; margin: auto; }
+    .title { text-align: center; }
+""") as demo:
+    with gr.Column(elem_classes="container"):
+        gr.Markdown("<h1 class='title'>ASR Demo with Editable Transcription</h1>")
+        with gr.Row():
+            audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Upload or Record Audio")
+            transcribe_button = gr.Button("Transcribe Audio", variant="primary")
+        with gr.Row():
+            original_text = gr.Textbox(label="Original Transcription", interactive=False, lines=5)
+            corrected_text = gr.Textbox(label="Corrected Transcription", interactive=True, lines=5)
+        # Hidden state to hold processing time
+        proc_time_state = gr.State()
+        with gr.Row():
+            save_button = gr.Button("Save Correction to Database", variant="primary")
+            save_status = gr.Textbox(label="Save Status", interactive=False)
+        with gr.Accordion("Download Options", open=False):
+            with gr.Row():
+                download_button = gr.Button("Download Results (ZIP)")
+                download_output = gr.File(label="Download ZIP")
+        
+        # Set up actions
+        transcribe_button.click(
+            fn=transcribe_both,
+            inputs=audio_input,
+            outputs=[original_text, corrected_text, proc_time_state]
+        )
+        
+        save_button.click(
+            fn=store_correction,
+            inputs=[original_text, corrected_text, audio_input, proc_time_state],
+            outputs=save_status
+        )
+        
+        download_button.click(
+            fn=prepare_download,
+            inputs=[audio_input, original_text, corrected_text],
+            outputs=download_output
+        )
 
 # Launch the demo
-demo.launch(share=True)
+demo.launch(share=True)