Update app.py

app.py

@@ -3,374 +3,281 @@ import os
 import requests
 import time
 
-# --- CONFIGURATION ---
+# ── Environment variables (set these in HF Space Secrets) ──────────────────────
 ASSEMBLYAI_API_KEY = os.environ.get("ASSEMBLYAI_API_KEY")
-HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# ============================================================
-# 1. SPEECH-TO-TEXT: AssemblyAI Integration
-# ============================================================
 
-def transcribe_audio_assemblyai(audio_file_path):
-    """Uses AssemblyAI's free tier (100 hours free)"""
-    if not ASSEMBLYAI_API_KEY:
-        return "❌ AssemblyAI API key not set. Add to Secrets."
-    
+# ══════════════════════════════════════════════════════════════════════════════
+# MODULE 1 — SPEECH-TO-TEXT (AssemblyAI)
+# Uploads audio to AssemblyAI, requests transcription, polls until complete.
+# Free tier: 100 hours/month. No local GPU needed.
+# ══════════════════════════════════════════════════════════════════════════════
+
+def transcribe_audio(audio_file_path):
     headers = {"authorization": ASSEMBLYAI_API_KEY}
-    
-    # Step 1: Upload audio
-    print("📤 Uploading to AssemblyAI...")
-    
-    def read_file(filename):
-        with open(filename, "rb") as f:
-            while True:
-                data = f.read(5242880)  # 5MB chunks
-                if not data:
-                    break
-                yield data
-    
-    upload_response = requests.post(
+
+    # Upload the audio file in 5 MB chunks
+    def read_file(path):
+        with open(path, "rb") as f:
+            while chunk := f.read(5_242_880):
+                yield chunk
+
+    upload_res = requests.post(
         "https://api.assemblyai.com/v2/upload",
         headers=headers,
         data=read_file(audio_file_path)
     )
-    
-    if upload_response.status_code != 200:
-        return f"❌ Upload failed: {upload_response.text}"
-    
-    audio_url = upload_response.json()["upload_url"]
-    print(f"✅ Uploaded: {audio_url}")
-    
-    # Step 2: Request transcription
-    json_data = {
-        "audio_url": audio_url,
-        "speech_models": ["universal-2"],  # Valid free tier model
-        "language_code": "en_us"
-    }
-    
-    transcript_response = requests.post(
+    if upload_res.status_code != 200:
+        return f"Upload failed: {upload_res.text}"
+
+    audio_url = upload_res.json()["upload_url"]
+
+    # Request transcription job
+    transcript_res = requests.post(
         "https://api.assemblyai.com/v2/transcript",
-        json=json_data,
-        headers=headers
+        headers=headers,
+        json={"audio_url": audio_url, "language_code": "en_us"}
     )
-    
-    if transcript_response.status_code != 200:
-        error_msg = transcript_response.json().get("error", "Unknown error")
-        return f"❌ Transcription request failed: {error_msg}"
-    
-    transcript_id = transcript_response.json()["id"]
-    print(f"📝 Transcript ID: {transcript_id}")
-    
-    # Step 3: Poll for results
-    polling_endpoint = f"https://api.assemblyai.com/v2/transcript/{transcript_id}"
-    
-    for attempt in range(30):  # Max 30 seconds
-        polling_response = requests.get(polling_endpoint, headers=headers)
-        polling_data = polling_response.json()
-        
-        status = polling_data["status"]
-        print(f"⏳ Status: {status}")
-        
-        if status == "completed":
-            print("✅ Transcription complete!")
-            return polling_data["text"]
-        elif status == "error":
-            return f"❌ Transcription error: {polling_data.get('error', 'Unknown')}"
-        
-        time.sleep(1)
-    
-    return "❌ Transcription timed out after 30 seconds"
-
-
-def transcribe_audio_placeholder(audio_file_path):
-    """Fallback when no API keys are available"""
-    return """
-Doctor: Hello, what brings you in today?
-Patient: I've had a cough for about two weeks. It gets worse at night and I feel really tired.
-Doctor: Any fever or shortness of breath?
-Patient: No fever, but I get winded climbing stairs.
-Doctor: I'm going to listen to your lungs. Take a deep breath. I can hear some mild wheezing on the right side.
-Patient: Is it serious?
-Doctor: It appears to be acute bronchitis. I'll prescribe an inhaler and recommend rest. Follow up in a week.
-Patient: Thank you, doctor.
-"""
+    if transcript_res.status_code != 200:
+        return f"Transcription request failed: {transcript_res.json().get('error', 'Unknown error')}"
 
+    transcript_id = transcript_res.json()["id"]
+    polling_url = f"https://api.assemblyai.com/v2/transcript/{transcript_id}"
 
-# ============================================================
-# 2. CLINICAL NOTE GENERATION: Rule-Based NLP
-# ============================================================
+    # Poll every 2 seconds until completed or failed (max 2 minutes)
+    for _ in range(60):
+        poll = requests.get(polling_url, headers=headers).json()
+        if poll["status"] == "completed":
+            return poll["text"]
+        if poll["status"] == "error":
+            return f"Transcription error: {poll.get('error', 'Unknown')}"
+        time.sleep(2)
 
-def generate_rule_based_note(transcript):
-    """Extracts clinical info using keyword matching and pattern recognition"""
-    t = transcript.lower()
-    
-    # Extract symptoms
-    symptoms = []
-    if "cough" in t: 
-        if "two week" in t or "2 week" in t:
-            symptoms.append("Cough (2 weeks duration)")
-        else:
-            symptoms.append("Cough")
-    if "fever" in t: 
-        symptoms.append("Fever")
-    if "tired" in t or "fatigue" in t: 
-        symptoms.append("Fatigue")
-    if "wheez" in t: 
-        symptoms.append("Wheezing")
-    if "breath" in t or "winded" in t: 
-        symptoms.append("Dyspnea on exertion")
-    if "night" in t and "cough" in t: 
-        symptoms.append("Nocturnal cough")
+    return "Timed out waiting for transcription. Please try again."
+
+
+def demo_transcript():
+    """Returns a sample transcript when no API key is configured."""
+    return (
+        "Doctor: Hello, what brings you in today?\n"
+        "Patient: I've had a cough for about two weeks. It gets worse at night and I feel really tired.\n"
+        "Doctor: Any fever or shortness of breath?\n"
+        "Patient: No fever, but I get winded climbing stairs.\n"
+        "Doctor: Let me listen to your lungs. I can hear some mild wheezing on the right side. "
+        "It appears to be acute bronchitis. I'll prescribe an inhaler and recommend rest. Follow up in a week.\n"
+        "Patient: Thank you, doctor."
+    )
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# MODULE 2 — CLINICAL NOTE GENERATION (Rule-Based NLP)
+# Extracts clinical entities from the transcript using keyword matching and
+# builds a structured SOAP note. This is the core NLP logic.
+# ══════════════════════════════════════════════════════════════════════════════
+
+def extract_symptoms(text):
+    """Identify symptoms mentioned in the transcript."""
+    t = text.lower()
+    found = []
+    if "cough" in t:
+        found.append("Cough — 2 weeks duration" if ("two week" in t or "2 week" in t) else "Cough")
+    if "fever" in t:
+        found.append("Fever")
+    if "tired" in t or "fatigue" in t:
+        found.append("Fatigue")
+    if "wheez" in t:
+        found.append("Wheezing")
+    if "winded" in t or "shortness of breath" in t or "dyspnea" in t:
+        found.append("Dyspnea on exertion")
+    if "night" in t and "cough" in t:
+        found.append("Nocturnal cough")
     if "chest" in t and "pain" in t:
-        symptoms.append("Chest pain")
+        found.append("Chest pain")
     if "headache" in t:
-        symptoms.append("Headache")
-    
-    # Determine diagnosis
+        found.append("Headache")
+    if "nausea" in t:
+        found.append("Nausea")
+    return found
+
+
+def determine_diagnosis(text):
+    """Map transcript keywords to a primary diagnosis."""
+    t = text.lower()
     if "bronchitis" in t:
-        diagnosis = "Acute Bronchitis"
-        confidence = "High"
-    elif "pneumonia" in t:
-        diagnosis = "Community-Acquired Pneumonia"
-        confidence = "Moderate"
-    elif "asthma" in t:
-        diagnosis = "Asthma Exacerbation"
-        confidence = "Moderate"
-    elif "covid" in t or "coronavirus" in t:
-        diagnosis = "COVID-19 Infection"
-        confidence = "Moderate"
-    elif "cough" in t and "wheez" in t:
-        diagnosis = "Acute Bronchitis with Reactive Airway Disease"
-        confidence = "Moderate"
-    elif "cough" in t and len(symptoms) >= 2:
-        diagnosis = "Upper Respiratory Infection"
-        confidence = "Moderate"
-    elif "cough" in t:
-        diagnosis = "Acute Cough, Etiology Pending"
-        confidence = "Low"
-    else:
-        diagnosis = "Pending Further Workup"
-        confidence = "Low"
-    
-    # Extract physical exam findings
-    exam_findings = []
-    if "wheez" in t: 
-        exam_findings.append("Mild expiratory wheezing on auscultation")
-    if "rhonchi" in t:
-        exam_findings.append("Rhonchi noted")
-    if "crackle" in t or "rale" in t:
-        exam_findings.append("Fine crackles at bases")
-    if "lung" in t and "clear" in t: 
-        exam_findings.append("Lungs clear bilaterally")
-    if not exam_findings: 
-        exam_findings.append("Unremarkable")
-    
-    # Build treatment plan
+        return "Acute Bronchitis"
+    if "pneumonia" in t:
+        return "Community-Acquired Pneumonia"
+    if "asthma" in t:
+        return "Asthma Exacerbation"
+    if "covid" in t or "coronavirus" in t:
+        return "COVID-19 Infection"
+    if "cough" in t and "wheez" in t:
+        return "Acute Bronchitis with Reactive Airway Disease"
+    if "cough" in t:
+        return "Upper Respiratory Infection"
+    return "Pending Further Workup"
+
+
+def build_plan(text):
+    """Construct a treatment plan based on clinical keywords."""
+    t = text.lower()
     plan = []
-    if "inhaler" in t or "wheez" in t: 
-        plan.append("- Albuterol HFA 90mcg, 2 puffs q4-6h PRN for wheezing")
+    if "inhaler" in t or "wheez" in t:
+        plan.append("Albuterol HFA 90 mcg — 2 puffs q4-6h PRN for wheezing")
     if "bronchitis" in t:
-        plan.append("- Supportive care (acute bronchitis typically viral, antibiotics not indicated)")
+        plan.append("Supportive care (acute bronchitis is typically viral; antibiotics not indicated)")
     if "antibiotic" in t:
-        plan.append("- Consider antibiotic therapy if bacterial infection suspected")
+        plan.append("Antibiotic therapy — consider if bacterial infection suspected")
     if "rest" in t or "tired" in t:
-        plan.append("- Recommend rest and increased fluid intake")
+        plan.append("Rest and increased fluid intake")
     if "cough" in t:
-        plan.append("- OTC dextromethorphan or guaifenesin for symptomatic cough relief")
-    
+        plan.append("OTC dextromethorphan or guaifenesin for symptomatic relief")
     if not plan:
-        plan.append("- Symptomatic management")
-    
-    plan.extend([
-        "- Avoid respiratory irritants and smoking",
-        "- Follow up in 7 days if symptoms persist or worsen",
-        "- Return to clinic sooner if fever develops or shortness of breath increases"
-    ])
-    
-    return f"""
-SUBJECTIVE:
-Chief Complaint: {symptoms[0] if symptoms else 'Not specified'}
-Associated Symptoms: {', '.join(symptoms[1:]) if len(symptoms) > 1 else 'None reported'}
-Duration: {'2 weeks' if 'two week' in t or '2 week' in t else 'Not specified'}
-Onset: {'Gradual' if 'week' in t else 'Not specified'}
-Severity: Moderate
-Aggravating Factors: {'Nighttime, exertion' if 'night' in t or 'breath' in t else 'None reported'}
-
-OBJECTIVE:
-Physical Exam: {', '.join(exam_findings)}
-Vital Signs: Temperature 98.6°F, HR 72, BP 118/76, RR 16, SpO2 97% on room air
-General: Alert, in no acute distress, well-appearing
-
-ASSESSMENT:
-Primary Diagnosis: {diagnosis}
-Clinical Confidence: {confidence}
-Differential Diagnoses: 
-- Viral Upper Respiratory Infection
-- Allergic Rhinitis with Post-nasal Drip
-- Asthma Exacerbation
-- GERD
-
-PLAN:
-{chr(10).join(plan)}
-
----
-**📋 Note**: This clinical note was generated using rule-based NLP extraction (keyword matching + pattern recognition) demonstrating the underlying logic used in production LLM fine-tuning. At Viscrow Health, the production pipeline used a fine-tuned Llama 3 8B model achieving 94% accuracy in clinical note generation.
-"""
+        plan.append("Symptomatic management")
+    # Standard advice added to every plan
+    plan.append("Avoid respiratory irritants")
+    plan.append("Follow up in 7 days, or sooner if symptoms worsen or fever develops")
+    return plan
+
+
+def generate_soap_note(transcript):
+    """
+    Orchestrates entity extraction and assembles the final SOAP note.
+    SOAP = Subjective / Objective / Assessment / Plan
+    """
+    t = transcript.lower()
+    symptoms  = extract_symptoms(transcript)
+    diagnosis = determine_diagnosis(transcript)
+    plan      = build_plan(transcript)
+
+    # Physical exam findings
+    findings = []
+    if "wheez" in t:
+        findings.append("Mild expiratory wheezing on auscultation")
+    if "rhonchi" in t:
+        findings.append("Rhonchi noted")
+    if "crackle" in t or "rale" in t:
+        findings.append("Fine crackles at lung bases")
+    if "lung" in t and "clear" in t:
+        findings.append("Lungs clear bilaterally")
+    if not findings:
+        findings.append("Unremarkable")
 
+    chief      = symptoms[0] if symptoms else "Not specified"
+    associated = ", ".join(symptoms[1:]) if len(symptoms) > 1 else "None reported"
+    duration   = "2 weeks" if ("two week" in t or "2 week" in t) else "Not specified"
+    agg        = "Nighttime, physical exertion" if ("night" in t or "winded" in t) else "None reported"
 
-def generate_clinical_note(transcript):
-    """Main clinical note generation function"""
-    if not transcript or len(transcript) < 20:
-        return "❌ Transcription too short. Please provide a longer audio file."
-    
-    if transcript.startswith("❌"):
-        return transcript
-    
-    # Use rule-based extraction (always works, no API needed)
-    return generate_rule_based_note(transcript)
+    plan_text = "\n".join(f"  - {item}" for item in plan)
 
+    return (
+        f"SUBJECTIVE\n"
+        f"{'─' * 42}\n"
+        f"Chief Complaint     : {chief}\n"
+        f"Associated Symptoms : {associated}\n"
+        f"Duration            : {duration}\n"
+        f"Aggravating Factors : {agg}\n"
+        f"Severity            : Moderate\n\n"
 
-# ============================================================
-# 3. MAIN PIPELINE
-# ============================================================
+        f"OBJECTIVE\n"
+        f"{'─' * 42}\n"
+        f"Vital Signs   : Temp 98.6F  HR 72  BP 118/76  RR 16  SpO2 97% RA\n"
+        f"General       : Alert, in no acute distress\n"
+        f"Physical Exam : {', '.join(findings)}\n\n"
+
+        f"ASSESSMENT\n"
+        f"{'─' * 42}\n"
+        f"Primary Diagnosis      : {diagnosis}\n"
+        f"Differential Diagnoses : Viral URI, Allergic rhinitis, Asthma exacerbation, GERD\n\n"
+
+        f"PLAN\n"
+        f"{'─' * 42}\n"
+        f"{plan_text}"
+    )
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# MAIN PIPELINE
+# Connects transcription -> note generation and handles missing API key gracefully.
+# ══════════════════════════════════════════════════════════════════════════════
 
 def process_encounter(audio):
-    """Main workflow: Audio → Transcription → SOAP Note"""
     if audio is None:
-        return "⚠️ Please upload an audio file.", ""
-    
-    print(f"\n{'='*60}")
-    print(f"🎤 Processing: {os.path.basename(audio)}")
-    print(f"📁 File size: {os.path.getsize(audio)} bytes")
-    
-    # Step 1: Transcribe audio
+        return "No audio provided. Please upload a file or use the microphone.", ""
+
+    # Step 1: Transcribe
     if ASSEMBLYAI_API_KEY:
-        print("🔑 Using AssemblyAI for transcription...")
-        transcript = transcribe_audio_assemblyai(audio)
+        transcript = transcribe_audio(audio)
     else:
-        print("⚠️ No AssemblyAI key - using sample transcript")
-        transcript = "⚠️ DEMO MODE - Add AssemblyAI API key to Secrets for live transcription\n\n"
-        transcript += transcribe_audio_placeholder(audio)
-    
-    print(f"📝 Transcript preview: {transcript[:150]}...")
-    
-    # Step 2: Generate clinical note
-    print("📋 Generating clinical note...")
-    note = generate_clinical_note(transcript)
-    
-    print(f"✅ Complete! Note length: {len(note)} chars")
-    print(f"{'='*60}\n")
-    
+        # No API key — use the built-in demo transcript so the app is still usable
+        transcript = "[Demo mode — add ASSEMBLYAI_API_KEY to Secrets for live transcription]\n\n" + demo_transcript()
+
+    if not transcript or len(transcript.strip()) < 20:
+        return transcript, "Transcript too short to generate a note."
+
+    # Step 2: Generate SOAP note
+    note = generate_soap_note(transcript)
+
     return transcript, note
 
 
-# ============================================================
-# 4. GRADIO USER INTERFACE
-# ============================================================
-
-demo = gr.Blocks(title="OpenScribe - Clinical AI Demo")
-
-with demo:
-    gr.Markdown("""
-    # 🏥 OpenScribe: AI Clinical Documentation
-    
-    ### **Educational Demonstration of the Viscrow Health Pipeline**
-    *Built by Arafat Anam Chowdhury*
-    
-    ---
-    
-    This tool replicates the **exact architecture** used in production for automated clinical documentation:
-    1. **Speech-to-Text**: AssemblyAI transcription (100 hours free tier)
-    2. **NLP Processing**: Rule-based clinical entity extraction
-    3. **Output**: Structured SOAP note ready for EHR integration
-    
-    *In production at Viscrow Health, this pipeline used fine-tuned Llama 3 8B for summarization, achieving 94% accuracy and reducing documentation time by 60%.*
+# ══════════════════════════════════════════════════════════════════════════════
+# GRADIO INTERFACE
+# ══════════════════════════════════════════════════════════════════════════════
+
+css = """
+    .gradio-container { max-width: 920px !important; margin: auto; }
+    footer { display: none !important; }
+    #app-title h1 { font-size: 1.35rem; font-weight: 600; margin: 0 0 0.2rem; }
+    #app-title p  { font-size: 0.82rem; color: #6b7280; margin: 0 0 1.25rem; padding-bottom: 1rem; border-bottom: 1px solid #e5e7eb; }
+    #hint { font-size: 0.8rem; color: #9ca3af; margin-top: 0.4rem; line-height: 1.5; }
+"""
+
+with gr.Blocks(
+    title="OpenScribe — AI Clinical Scribe",
+    theme=gr.themes.Soft(primary_hue="slate", neutral_hue="slate"),
+    css=css
+) as demo:
+
+    gr.HTML("""
+        <div id="app-title">
+            <h1>OpenScribe</h1>
+            <p>AI clinical scribe &mdash; upload a doctor&ndash;patient recording to generate a structured SOAP note.</p>
+        </div>
     """)
-    
-    with gr.Row():
-        with gr.Column(scale=1):
+
+    with gr.Row(equal_height=False):
+
+        # Left column — input
+        with gr.Column(scale=1, min_width=240):
             audio_input = gr.Audio(
                 type="filepath",
-                label="📁 Upload Medical Conversation",
+                label="Recording",
                 sources=["upload", "microphone"]
             )
-            
-            run_btn = gr.Button(
-                "📋 Generate Clinical Note", 
-                variant="primary", 
-                size="lg"
-            )
-            
-            # Status indicators
-            with gr.Group():
-                gr.Markdown("### 🔧 System Status")
-                if ASSEMBLYAI_API_KEY:
-                    gr.Markdown("✅ **AssemblyAI:** Connected")
-                else:
-                    gr.Markdown("⚠️ **AssemblyAI:** Not configured (demo mode)")
-                
-                gr.Markdown("✅ **NLP Engine:** Rule-Based Extraction (Active)")
-            
-            gr.Markdown("""
-            ---
-            ### 📁 Sample Files
-            
-            **Test Audio Files:**
-            - [Medical WAV Sample](https://www.voiptroubleshooter.com/open_speech/american/OSR_us_000_0010_8k.wav)
-            
-            **Or record your own conversation:**
-            *"Hi, what brings you in? - I've had this cough for two weeks. - Any fever? - No. - Let me listen... I hear wheezing. - It's bronchitis."*
-            """)
-        
+            run_btn = gr.Button("Generate note", variant="primary")
+            gr.HTML('<p id="hint">Supports MP3, WAV, M4A.<br>No file? Record yourself reading a short mock encounter.</p>')
+
+        # Right column — outputs
         with gr.Column(scale=2):
-            transcript_output = gr.Textbox(
-                label="📝 Step 1: Transcription",
+            transcript_out = gr.Textbox(
+                label="Transcript",
                 lines=6,
-                placeholder="Transcribed conversation will appear here..."
+                placeholder="Transcribed conversation will appear here...",
+                show_copy_button=True
             )
-
-            note_output = gr.Textbox(
-                label="📋 Step 2: Generated SOAP Note",
-                lines=20,
-                placeholder="Clinical documentation will appear here..."
+            note_out = gr.Textbox(
+                label="SOAP Note",
+                lines=18,
+                placeholder="Generated clinical note will appear here...",
+                show_copy_button=True
             )
-    
+
     run_btn.click(
         fn=process_encounter,
         inputs=audio_input,
-        outputs=[transcript_output, note_output]
+        outputs=[transcript_out, note_out]
     )
-    
-    gr.Markdown("""
-    ---
-    
-    ### 🔬 Technical Implementation
-    
-    | Component | This Demo | Production (Viscrow Health) |
-    |-----------|-----------|------------------------------|
-    | **Speech-to-Text** | AssemblyAI Universal-2 | Azure Speech Services / Whisper |
-    | **Entity Extraction** | Rule-Based NLP (Keyword + Pattern) | Fine-tuned Llama 3 8B |
-    | **Output Format** | SOAP Note | SOAP Note + Billing Codes |
-    | **Error Handling** | Multi-tier Fallback | Validation Pipeline |
-    
-    ### 📊 Key Achievements (Viscrow Health)
-    
-    - ✅ Integrated speech-to-text and LLM summarization pipelines
-    - ✅ Designed tools linking clinician notes to billing details
-    - ✅ Evaluated AI outputs and reduced common errors
-    - ✅ Built React frontend components for clinical dashboards
-    
-    ---
-    
-    **⚠️ Educational Disclaimer**: This is a portfolio demonstration. Not for real clinical use.
-    
-    [GitHub](https://github.com/arafatanam) | [LinkedIn](https://www.linkedin.com/in/arafat-anam-chowdhury) | [Hugging Face](https://huggingface.co/arafatanam)
-    """)
-
-
-# ============================================================
-# 5. LAUNCH
-# ============================================================
 
 if __name__ == "__main__":
-    demo.launch(theme=gr.themes.Soft())
+    demo.launch()