Update app.py

app.py

@@ -10,7 +10,8 @@ from datasets import load_dataset, Audio
 from gradio_client import Client
 from stats_data import get_indomain_breakdown, get_experimental_summary, SPEAKER_META
 
-# 1. Setup Local Whisper Baseline (English, Strict Generation)
+# 1. Initialize Baseline ASR (Strict English, Repetition Penalty 3.0)
+print("Loading Whisper Tiny...")
 whisper_asr = pipeline(
     "automatic-speech-recognition", 
     model="openai/whisper-tiny",
@@ -20,119 +21,123 @@ whisper_asr = pipeline(
 HF_TOKEN = os.getenv("HF_TOKEN")
 PRIVATE_BACKEND_URL = "st192011/Torgo-DSR-Private"
 
-def normalize_text(text):
+def normalize(text):
     if not text: return ""
     return re.sub(r'[^\w\s]', '', text).lower().strip()
 
-# --- Shared Processing Logic ---
-
-def process_audio_file(audio_path):
-    """Ensures any input audio is formatted correctly for ASR systems (16kHz Mono)."""
-    y, sr = librosa.load(audio_path, sr=16000)
-    fixed_path = "processed_audio.wav"
-    sf.write(fixed_path, y, sr)
-    return fixed_path
+# --- Logic: Data Loading ---
+def get_sample_logic(speaker_id):
+    try:
+        if "UA" in speaker_id:
+            # UA-Speech Access (Direct pull for F02)
+            dataset = load_dataset("resproj007/uaspeech_female", split="train", streaming=True)
+            dataset = dataset.cast_column("audio", Audio(decode=False))
+            # UA is small, skip slightly for variety
+            sample = next(iter(dataset.skip(random.randint(0, 30))))
+            gt_text = sample.get('text') or sample.get('transcription') or sample.get('sentence')
+        else:
+            # Torgo Access (Manual filtering as per Colab fix)
+            dataset = load_dataset("abnerh/TORGO-database", split="train", streaming=True)
+            dataset = dataset.cast_column("audio", Audio(decode=False))
+            
+            def filter_spk(x):
+                sid = str(x.get('speaker_id', '')).upper()
+                if not sid or sid == "NONE":
+                    sid = os.path.basename(x['audio']['path']).split('_')[0].upper()
+                return sid == speaker_id
+            
+            speaker_ds = dataset.filter(filter_spk)
+            sample = next(iter(speaker_ds.shuffle(buffer_size=10)))
+            gt_text = sample.get('transcription') or sample.get('text')
+
+        # Decode Bytes manually to bypass torchcodec errors
+        audio_bytes = sample['audio']['bytes']
+        audio_data, sr = librosa.load(io.BytesIO(audio_bytes), sr=16000)
+        temp_path = "sample.wav"
+        sf.write(temp_path, audio_data, sr)
+        
+        return temp_path, gt_text.lower().strip(), SPEAKER_META[speaker_id]
+    except Exception as e:
+        return None, f"Dataset Error: {e}", {}
 
-def run_whisper_logic(audio_path):
+# --- Logic: Model Steps ---
+def run_whisper_step(audio_path):
     if not audio_path: return "No audio loaded", ""
-    formatted_path = process_audio_file(audio_path)
-    result = whisper_asr(formatted_path)
+    result = whisper_asr(audio_path)
     raw_w = result["text"]
-    norm_w = normalize_text(raw_w)
+    norm_w = normalize(raw_w)
     return raw_w, norm_w
 
-def run_reconstruction_logic(audio_path, norm_whisper):
-    if not audio_path or not norm_whisper: return "Run Whisper step first."
+def run_model_step(audio_path, norm_whisper):
+    if not audio_path or not norm_whisper: return "Complete Steps 1 & 2 first."
     try:
+        # Call the private space for the 5K Gemma Model prediction
         client = Client(PRIVATE_BACKEND_URL, hf_token=HF_TOKEN)
-        # Private backend handles Wav2Vec, Allosaurus, and Gemma 3 arbitration
         prediction = client.predict(audio_path, norm_whisper, api_name="/predict_dsr")
         return prediction
     except Exception as e:
-        return f"Backend Offline. Error: {e}"
-
-# --- Channel 1: Dataset Loader ---
-def get_dataset_sample(speaker_id):
-    try:
-        if speaker_id == "F02":
-            ds = load_dataset("resproj007/uaspeech_female", split="train", streaming=True)
-            ds = ds.cast_column("audio", Audio(decode=False))
-            sample = next(iter(ds.skip(random.randint(0, 50))))
-            gt_text = sample.get('text') or sample.get('transcription') or "Unknown"
-        else:
-            ds = load_dataset("abnerh/TORGO-database", split="train", streaming=True)
-            ds = ds.cast_column("audio", Audio(decode=False))
-            indices = {'M05': 10978, 'M02': 11565, 'M04': 12337, 'M01': 13003, 'F01': 13746, 'M03': 13982, 'F04': 14792, 'F03': 15465}
-            start_idx = indices.get(speaker_id, 0)
-            sample = next(iter(ds.skip(start_idx + random.randint(0, 10))))
-            gt_text = sample.get('transcription') or sample.get('text') or "Unknown"
-
-        audio_data, sr = librosa.load(io.BytesIO(sample['audio']['bytes']), sr=16000)
-        temp_path = f"sample_{speaker_id}.wav"
-        sf.write(temp_path, audio_data, sr)
-        return temp_path, gt_text.lower().strip(), SPEAKER_META.get(speaker_id, {})
-    except Exception as e:
-        return None, f"Dataset Error: {e}", {}
+        return f"Backend Offline. Research Details: {e}"
 
-# --- UI Layout ---
+# --- UI Construction ---
 with gr.Blocks(theme=gr.themes.Soft(), title="Torgo DSR Lab") as demo:
     gr.Markdown("# ⚗️ Torgo DSR Lab")
-    gr.Markdown("ASR Correction and Reconstruction Layer for Torgo and UA-Speech.")
+    gr.Markdown("Neural Reconstruction for Severe Dysarthria benchmarked on Torgo and UA-Speech.")
     
-    # States for audio paths
-    lab_audio_state = gr.State("")
-    user_audio_state = gr.State("")
+    current_audio_path = gr.State("")
 
-    with gr.Tab("🔬 Research Samples"):
-        gr.Markdown("Select clinical samples from the Torgo or UA-Speech datasets.")
+    with gr.Tab("🔬 Laboratory"):
         with gr.Row():
             with gr.Column(scale=1):
+                gr.Markdown("### Step 1: Load Sample")
                 speaker_input = gr.Dropdown(sorted(list(SPEAKER_META.keys())), label="Speaker ID", value="F01")
-                load_btn = gr.Button("Load Sample Data")
-                meta_display = gr.JSON(label="Sample Metadata")
+                load_btn = gr.Button("Load Data")
+                meta_display = gr.JSON(label="Speaker Meta")
                 gt_box = gr.Textbox(label="Ground Truth")
-            
-            with gr.Column(scale=2):
-                whisper_btn_lab = gr.Button("1. Generate Whisper Baseline")
-                w_raw_lab = gr.Textbox(label="Whisper Raw")
-                w_norm_lab = gr.Textbox(label="Whisper Normalized")
-                
-                gr.Markdown("---")
-                model_btn_lab = gr.Button("2. Run Neural Reconstruction", variant="primary")
-                final_out_lab = gr.Textbox(label="DSR Lab Prediction")
+                # Added visible audio for user verification
+                audio_preview = gr.Audio(label="Audio Preview", type="filepath")
 
-    with gr.Tab("🎤 Personal Test"):
-        gr.Markdown("Record or upload your own audio to test the reconstruction layer.")
-        with gr.Row():
-            with gr.Column(scale=1):
-                user_audio_input = gr.Audio(sources=["microphone", "upload"], type="filepath", label="User Audio")
-                process_user_btn = gr.Button("Prepare Audio")
-            
             with gr.Column(scale=2):
-                whisper_btn_user = gr.Button("1. Generate Whisper Baseline")
-                w_raw_user = gr.Textbox(label="Whisper Raw")
-                w_norm_user = gr.Textbox(label="Whisper Normalized")
+                gr.Markdown("### Step 2: ASR Baseline")
+                whisper_btn = gr.Button("Run Whisper Tiny")
+                w_raw = gr.Textbox(label="Whisper Raw")
+                w_norm = gr.Textbox(label="Whisper Normalized")
                 
                 gr.Markdown("---")
-                model_btn_user = gr.Button("2. Run Neural Reconstruction", variant="primary")
-                final_out_user = gr.Textbox(label="DSR Lab Prediction")
+                gr.Markdown("### Step 3: Neural Reconstruction")
+                model_btn = gr.Button("Run Our Model", variant="primary")
+                final_out = gr.Textbox(label="DSR Lab Prediction")
 
     with gr.Tab("📊 Research Statistics"):
-        gr.Markdown("# 🔬 Scientific Evaluation")
-        gr.Markdown("**Metric:** Exact Match Accuracy on normalized text (lowercase, no punctuation).")
-        gr.Markdown("## 1. Torgo In-Domain Breakdown")
+        gr.Markdown("# 🔬 Performance Evaluation")
+        
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("""
+                ### 📏 Metric: Exact Match Accuracy
+                Accuracy is the percentage of samples where the **normalized prediction** (lowercase, no punctuation) matches the **ground truth**.
+                """)
+            
+            with gr.Column():
+                gr.Markdown("""
+                ### 🧪 Model Definitions
+                * **5K Pure Model:** Trained on real articulatory distortions. Optimized for phonetic fidelity.
+                * **10K Triple-Mix Model:** Includes synthetic data and anchors; utilized for generalization testing.
+                """)
+
+        gr.Markdown("## 1. Torgo In-Domain Analysis")
         gr.DataFrame(get_indomain_breakdown())
-        gr.Markdown("## 2. Experimental Milestone Summary")
+        
+        gr.Markdown("## 2. Experimental Summary")
         gr.DataFrame(get_experimental_summary())
 
-    # --- Events: Research Tab ---
-    load_btn.click(get_dataset_sample, inputs=speaker_input, outputs=[lab_audio_state, gt_box, meta_display])
-    whisper_btn_lab.click(run_whisper_logic, inputs=lab_audio_state, outputs=[w_raw_lab, w_norm_lab])
-    model_btn_lab.click(run_reconstruction_logic, inputs=[lab_audio_state, w_norm_lab], outputs=final_out_lab)
+    # Event Mapping
+    load_btn.click(
+        get_sample_logic, 
+        inputs=speaker_input, 
+        outputs=[current_audio_path, gt_box, meta_display]
+    ).then(lambda x: x, inputs=current_audio_path, outputs=audio_preview)
 
-    # --- Events: Personal Tab ---
-    process_user_btn.click(lambda x: x, inputs=user_audio_input, outputs=user_audio_state)
-    whisper_btn_user.click(run_whisper_logic, inputs=user_audio_state, outputs=[w_raw_user, w_norm_user])
-    model_btn_user.click(run_reconstruction_logic, inputs=[user_audio_state, w_norm_user], outputs=final_out_user)
+    whisper_btn.click(run_whisper_step, inputs=current_audio_path, outputs=[w_raw, w_norm])
+    model_btn.click(run_model_step, inputs=[current_audio_path, w_norm], outputs=final_out)
 
 demo.launch()