Update for chinese and fix ui (#5)

- Update for chinese and fix ui (617bfb22d00ff3c86fb7cfabba3c3c382815efd5)


Co-authored-by: Hunter S <hunterschep@users.noreply.huggingface.co>

app.py

@@ -12,7 +12,7 @@ import tempfile
 
 # Initialize Firebase
 firebase_config = json.loads(os.environ.get('firebase_creds'))
-cred = credentials.Certificate(firebase_config)  # Your Firebase JSON key file
+cred = credentials.Certificate(firebase_config)
 firebase_admin.initialize_app(cred)
 db = firestore.client()
 
@@ -21,17 +21,54 @@ MODEL_NAME = "eleferrand/xlsr53_Amis"
 processor = Wav2Vec2Processor.from_pretrained(MODEL_NAME)
 model = AutoModelForCTC.from_pretrained(MODEL_NAME)
 
+# Language configuration
+LANGUAGE = {
+    "en": {
+        "title": "ASR Demo with Editable Transcription",
+        "step1": "Step 1: Audio Upload & Transcription",
+        "audio_input": "Audio Input",
+        "transcribe_btn": "Transcribe Audio",
+        "step2": "Step 2: Review & Edit Transcription",
+        "original_text": "Original Transcription",
+        "corrected_text": "Corrected Transcription",
+        "transcription_placeholder": "Transcription will appear here...",
+        "step3": "Step 3: User Information",
+        "age_label": "Age",
+        "native_speaker": "Native Amis Speaker",
+        "step4": "Step 4: Save & Download",
+        "save_btn": "Save Correction to Database",
+        "save_status": "Save Status",
+        "download_btn": "Download Results (ZIP)",
+        "status_placeholder": "Status messages will appear here...",
+        "toggle_lang": "中文/English"
+    },
+    "zh": {
+        "title": "可編輯轉寫的語音辨識演示",
+        "step1": "步驟一: 音頻上傳與轉寫",
+        "audio_input": "音頻輸入",
+        "transcribe_btn": "開始轉寫",
+        "step2": "步驟二: 校對與編輯轉寫結果",
+        "original_text": "原始轉寫結果",
+        "corrected_text": "校正後文本",
+        "transcription_placeholder": "轉寫結果將顯示在此處...",
+        "step3": "步驟三: 用戶資訊",
+        "age_label": "年齡",
+        "native_speaker": "阿美族母語者",
+        "step4": "步驟四: 保存與下載",
+        "save_btn": "保存校正結果至數據庫",
+        "save_status": "保存狀態",
+        "download_btn": "下載結果(ZIP壓縮檔)",
+        "status_placeholder": "狀態訊息將顯示在此處...",
+        "toggle_lang": "English/中文"
+    }
+}
+
+current_lang = gr.State(value="en")
+
 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)
         input_values = processor(audio, sampling_rate=16000, return_tensors="pt").input_values
-
-        # Get model predictions and decode to text
         with torch.no_grad():
             logits = model(input_values).logits
         predicted_ids = torch.argmax(logits, dim=-1)
@@ -41,22 +78,12 @@ def transcribe(audio_file):
         return f"Error processing file: {e}"
 
 def transcribe_both(audio_file):
-    """
-    Transcribes the audio and returns:
-      - the original transcription (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)
     processing_time = (datetime.now() - start_time).total_seconds()
     return transcription, transcription, processing_time
 
 def store_correction(original_transcription, corrected_transcription, audio_file, processing_time, age, native_speaker):
-    """
-    Stores the transcriptions and additional metadata (including user info and audio details)
-    in Firestore as a single document.
-    """
     try:
         audio_metadata = {}
         if audio_file and os.path.exists(audio_file):
@@ -70,7 +97,7 @@ def store_correction(original_transcription, corrected_transcription, audio_file
             'timestamp': datetime.now().isoformat(),
             'processing_time': processing_time,
             'audio_metadata': audio_metadata,
-            'audio_url': None,  # Placeholder if you decide to store an URL later
+            'audio_url': None,
             'model_name': MODEL_NAME,
             'user_info': {
                 'native_amis_speaker': native_speaker,
@@ -83,34 +110,18 @@ def store_correction(original_transcription, corrected_transcription, audio_file
         return f"Error saving correction: {e}"
 
 def prepare_download(audio_file, original_transcription, corrected_transcription):
-    """
-    Prepares a ZIP file containing:
-      - 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
-
-    # Create a temporary file to avoid filename conflicts
     tmp_zip = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
     tmp_zip.close()
     with zipfile.ZipFile(tmp_zip.name, "w") as zf:
-        # 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)
-
-        # Write and add the original transcription
         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)
-
-        # Write and add the corrected transcription
         corr_txt = "corrected_transcription.txt"
         with open(corr_txt, "w", encoding="utf-8") as f:
             f.write(corrected_transcription)
@@ -118,96 +129,117 @@ def prepare_download(audio_file, original_transcription, corrected_transcription
         os.remove(corr_txt)
     return tmp_zip.name
 
-# Build the Gradio Blocks interface with improved styling
+def toggle_language(lang):
+    new_lang = "zh" if lang == "en" else "en"
+    lang_dict = LANGUAGE[new_lang]
+    return [
+        gr.Markdown.update(value=f"<h1 class='header'>{lang_dict['title']}</h1>"),
+        gr.Markdown.update(value=f"### {lang_dict['step1']}"),
+        gr.Audio.update(label=lang_dict['audio_input']),
+        gr.Button.update(value=lang_dict['transcribe_btn']),
+        gr.Markdown.update(value=f"### {lang_dict['step2']}"),
+        gr.Textbox.update(label=lang_dict['original_text'], placeholder=lang_dict['transcription_placeholder']),
+        gr.Textbox.update(label=lang_dict['corrected_text'], placeholder=lang_dict['transcription_placeholder']),
+        gr.Markdown.update(value=f"### {lang_dict['step3']}"),
+        gr.Slider.update(label=lang_dict['age_label']),
+        gr.Checkbox.update(label=lang_dict['native_speaker']),
+        gr.Markdown.update(value=f"### {lang_dict['step4']}"),
+        gr.Button.update(value=lang_dict['save_btn']),
+        gr.Textbox.update(label=lang_dict['save_status'], placeholder=lang_dict['status_placeholder']),
+        gr.Button.update(value=lang_dict['download_btn']),
+        gr.File.update(label=lang_dict['download_btn']),
+        gr.Button.update(value=lang_dict['toggle_lang']),
+        new_lang
+    ]
+
 with gr.Blocks(css="""
-    .container { 
-        max-width: 800px; 
-        margin: auto; 
-        padding: 20px; 
-        font-family: Arial, sans-serif;
-    }
-    .header { 
-        text-align: center; 
-        margin-bottom: 30px; 
-    }
-    .section { 
-        margin-bottom: 30px; 
-        padding: 15px; 
-        border: 1px solid #ddd; 
-        border-radius: 8px; 
-        background-color: #f9f9f9;
-    }
-    .section h3 {
-        margin-top: 0;
-        margin-bottom: 15px;
-        text-align: center;
-    }
-    .button-row {
-        display: flex;
-        justify-content: center;
-        gap: 10px;
-        flex-wrap: wrap;
-    }
-    @media (max-width: 600px) {
-      .gradio-row { 
-          flex-direction: column; 
-      }
-    }
+    .container { max-width: 800px; margin: auto; padding: 20px; font-family: Arial, sans-serif; }
+    .header { text-align: center; margin-bottom: 30px; }
+    .section { margin-bottom: 30px; padding: 15px; border: 1px solid #ddd; border-radius: 8px; background-color: #f9f9f9; }
+    .section h3 { margin-top: 0; margin-bottom: 15px; text-align: center; }
+    .button-row { display: flex; justify-content: center; gap: 10px; flex-wrap: wrap; }
+    .lang-toggle { position: absolute; top: 20px; right: 20px; }
+    @media (max-width: 600px) { .gradio-row { flex-direction: column; } }
 """) as demo:
+    current_lang.render()
+    
     with gr.Column(elem_classes="container"):
-        gr.Markdown("<h1 class='header'>ASR Demo with Editable Transcription</h1>")
+        with gr.Row():
+            title_md = gr.Markdown(elem_classes="header")
+            lang_btn = gr.Button(LANGUAGE['en']['toggle_lang'], elem_classes="lang-toggle")
         
-        # Step 1: Audio Upload & Transcription
-        with gr.Row(elem_classes="section"):
-            gr.Markdown("### Step 1: Audio Upload & Transcription")
-            with gr.Row(elem_classes="gradio-row"):
-                audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath", label="Audio Input")
-                transcribe_button = gr.Button("Transcribe Audio", variant="primary")
+        # Step 1
+        with gr.Column(elem_classes="section"):
+            step1_md = gr.Markdown()
+            with gr.Row():
+                audio_input = gr.Audio(sources=["upload", "microphone"], type="filepath")
+                transcribe_button = gr.Button(variant="primary")
             proc_time_state = gr.State()
         
-        # Step 2: Review & Edit Transcription
-        with gr.Row(elem_classes="section"):
-            gr.Markdown("### Step 2: Review & Edit Transcription")
-            with gr.Row(elem_classes="gradio-row"):
-                original_text = gr.Textbox(label="Original Transcription", interactive=False, lines=5, placeholder="Transcription will appear here...")
-                corrected_text = gr.Textbox(label="Corrected Transcription", interactive=True, lines=5, placeholder="Edit transcription here...")
+        # Step 2
+        with gr.Column(elem_classes="section"):
+            step2_md = gr.Markdown()
+            with gr.Row():
+                original_text = gr.Textbox(interactive=False, lines=5)
+                corrected_text = gr.Textbox(interactive=True, lines=5)
         
-        # Step 3: User Information
-        with gr.Row(elem_classes="section"):
-            gr.Markdown("### Step 3: User Information")
-            with gr.Row(elem_classes="gradio-row"):
-                age_input = gr.Slider(minimum=0, maximum=100, step=1, label="Age", value=25)
-                native_speaker_input = gr.Checkbox(label="Native Amis Speaker", value=True)
+        # Step 3
+        with gr.Column(elem_classes="section"):
+            step3_md = gr.Markdown()
+            with gr.Row():
+                age_input = gr.Slider(minimum=0, maximum=100, step=1, value=25)
+                native_speaker_input = gr.Checkbox(value=True)
         
-        # Step 4: Save & Download
-        with gr.Row(elem_classes="section"):
-            gr.Markdown("### Step 4: Save & Download")
+        # Step 4
+        with gr.Column(elem_classes="section"):
+            step4_md = gr.Markdown()
             with gr.Row(elem_classes="button-row"):
-                save_button = gr.Button("Save Correction to Database", variant="primary")
-                save_status = gr.Textbox(label="Save Status", interactive=False, placeholder="Status messages will appear here...")
+                save_button = gr.Button(variant="primary")
+                save_status = gr.Textbox(interactive=False)
             with gr.Row(elem_classes="button-row"):
-                download_button = gr.Button("Download Results (ZIP)")
-                download_output = gr.File(label="Download ZIP")
-        
-        # UI 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, age_input, native_speaker_input],
-            outputs=save_status
-        )
-        
-        download_button.click(
-            fn=prepare_download,
-            inputs=[audio_input, original_text, corrected_text],
-            outputs=download_output
-        )
-
+                download_button = gr.Button()
+                download_output = gr.File()
+    
+    lang_btn.click(
+        toggle_language,
+        inputs=current_lang,
+        outputs=[
+            title_md, step1_md, audio_input, transcribe_button,
+            step2_md, original_text, corrected_text, step3_md,
+            age_input, native_speaker_input, step4_md, save_button,
+            save_status, download_button, download_output, lang_btn,
+            current_lang
+        ]
+    )
+    
+    transcribe_button.click(
+        transcribe_both,
+        inputs=audio_input,
+        outputs=[original_text, corrected_text, proc_time_state]
+    )
+    
+    save_button.click(
+        store_correction,
+        inputs=[original_text, corrected_text, audio_input, proc_time_state, age_input, native_speaker_input],
+        outputs=save_status
+    )
+    
+    download_button.click(
+        prepare_download,
+        inputs=[audio_input, original_text, corrected_text],
+        outputs=download_output
+    )
+    
+    demo.load(
+        toggle_language,
+        inputs=current_lang,
+        outputs=[
+            title_md, step1_md, audio_input, transcribe_button,
+            step2_md, original_text, corrected_text, step3_md,
+            age_input, native_speaker_input, step4_md, save_button,
+            save_status, download_button, download_output, lang_btn,
+            current_lang
+        ]
+    )
 
-# Launch the demo
 demo.launch(share=True)