Update app.py

app.py

@@ -5,7 +5,6 @@ import traceback
 import pandas as pd
 import torch
 import gradio as gr
-import gc
 from transformers import (
     logging,
     AutoProcessor,
@@ -13,6 +12,7 @@ from transformers import (
     AutoModelForImageTextToText
 )
 from sklearn.model_selection import train_test_split
+import gc
 
 # ─── Silence irrelevant warnings ───────────────────────────────────────────────
 logging.set_verbosity_error()
@@ -35,6 +35,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 def generate_and_export():
     try:
         # 1) Lazy‑load the full FP16 model
+        print("Loading model...")
         model = AutoModelForImageTextToText.from_pretrained(
             MODEL_ID,
             trust_remote_code=True,
@@ -43,25 +44,37 @@ def generate_and_export():
             device_map="auto"
         )
         device = next(model.parameters()).device
+        print(f"Model loaded on device: {device}")
 
-        # 2) Text→SOAP helper
-        def to_soap(text: str) -> str:
+        # 2) Text→Doctor Note helper
+        def generate_doctor_note() -> str:
+            prompt = """Generate a realistic, concise doctor's progress note for a single patient encounter. 
+            Include patient symptoms, physical examination findings, and clinical observations. 
+            Keep it brief and medical in nature.
+            
+            Example format:
+            Patient presents with [symptoms]. Physical exam reveals [findings]. [Additional observations].
+            
+            Doctor's Note:"""
+            
             inputs = processor.apply_chat_template(
                 [
-                    {"role":"system","content":[{"type":"text","text":"You are a medical AI assistant."}]},
-                    {"role":"user",  "content":[{"type":"text","text":text}]}
+                    {"role": "system", "content": [{"type": "text", "text": "You are a medical assistant generating realistic patient encounter notes."}]},
+                    {"role": "user", "content": [{"type": "text", "text": prompt}]}
                 ],
                 add_generation_prompt=True,
                 tokenize=True,
                 return_tensors="pt",
                 return_dict=True
             ).to(device)
+            
             out = model.generate(
                 **inputs,
-                max_new_tokens=400,
+                max_new_tokens=150,
                 do_sample=True,
-                top_p=0.95,
-                temperature=0.1,
+                top_p=0.9,
+                temperature=0.8,
+                repetition_penalty=1.2,
                 pad_token_id=processor.tokenizer.eos_token_id,
                 use_cache=False
             )
@@ -70,82 +83,203 @@ def generate_and_export():
                 out[:, prompt_len:], skip_special_tokens=True
             )[0].strip()
 
-        # 3) Generate 20 doc notes + ground truths
-        docs, gts = [], []
-        for i in range(1, 21):
-            doc = to_soap("Generate a realistic, concise doctor's progress note for a single patient encounter.")
-            docs.append(doc)
-            gts.append(to_soap(doc))
-            if i % 5 == 0:
-                torch.cuda.empty_cache()
+        # 3) Doctor Note→SOAP helper
+        def convert_to_soap(doctor_note: str) -> str:
+            prompt = f"""Convert the following medical note into proper SOAP format.
 
-        # 4) Split into 15 train / 5 test
-        df = pd.DataFrame({"doc_note": docs, "ground_truth_soap": gts})
-        train_df, test_df = train_test_split(df, test_size=5, random_state=42)
+Medical Note: {doctor_note}
 
-        os.makedirs("outputs", exist_ok=True)
+Please structure your response exactly as follows:
 
-        # 5) Inference on train split → outputs/inference.tsv
-        train_preds = [to_soap(d) for d in train_df["doc_note"]]
-        inf = train_df.reset_index(drop=True).copy()
-        inf["id"]             = inf.index + 1
-        inf["predicted_soap"] = train_preds
-        inf[["id","ground_truth_soap","predicted_soap"]].to_csv(
-            "outputs/inference.tsv", sep="\t", index=False
-        )
+SUBJECTIVE:
+[Patient's reported symptoms, complaints, and history]
 
-        # 6) Inference on test split → outputs/eval.csv
-        test_preds = [to_soap(d) for d in test_df["doc_note"]]
-        pd.DataFrame({
-            "id":             range(1, len(test_preds) + 1),
-            "predicted_soap": test_preds
-        }).to_csv("outputs/eval.csv", index=False)
+OBJECTIVE:
+[Physical exam findings, vital signs, observable data]
 
-        # 7) Return status + file paths for download
-        return (
-            "✅ Done with 20 notes (15 train / 5 test)!",
-            "outputs/inference.tsv",
-            "outputs/eval.csv"
-        )
-        # Add this to your generate_and_export function
+ASSESSMENT:
+[Clinical diagnosis, differential diagnosis, or impression]
+
+PLAN:
+[Treatment plan, medications, follow-up instructions]
+
+SOAP Note:"""
+            
+            inputs = processor.apply_chat_template(
+                [
+                    {"role": "system", "content": [{"type": "text", "text": "You are a medical documentation assistant. Convert medical notes into structured SOAP format. Be concise and clinical."}]},
+                    {"role": "user", "content": [{"type": "text", "text": prompt}]}
+                ],
+                add_generation_prompt=True,
+                tokenize=True,
+                return_tensors="pt",
+                return_dict=True
+            ).to(device)
+            
+            out = model.generate(
+                **inputs,
+                max_new_tokens=250,
+                do_sample=True,
+                top_p=0.9,
+                temperature=0.7,
+                repetition_penalty=1.3,
+                pad_token_id=processor.tokenizer.eos_token_id,
+                use_cache=False
+            )
+            prompt_len = inputs["input_ids"].shape[-1]
+            return processor.batch_decode(
+                out[:, prompt_len:], skip_special_tokens=True
+            )[0].strip()
 
-        # ... existing code ...
+        # 4) Generate 20 doctor notes + convert to SOAP
+        print("Generating doctor notes and SOAP conversions...")
+        docs, soaps = [], []
         
-        # Better memory management
         for i in range(1, 21):
-            doc = to_soap("Generate a realistic, concise doctor's progress note...")
-            docs.append(doc)
-            gts.append(to_soap(doc))
+            print(f"Generating note {i}/20...")
+            
+            # Generate doctor note
+            doctor_note = generate_doctor_note()
+            docs.append(doctor_note)
+            
+            # Convert to SOAP
+            soap_note = convert_to_soap(doctor_note)
+            soaps.append(soap_note)
             
-            # More aggressive cleanup every 3 iterations
+            # Memory cleanup every 3 iterations
             if i % 3 == 0:
                 torch.cuda.empty_cache()
                 gc.collect()
+                print(f"Memory cleaned after note {i}")
+
+        print("All notes generated successfully!")
+
+        # 5) Split into 15 train / 5 test
+        df = pd.DataFrame({"doctor_note": docs, "soap_note": soaps})
+        train_df, test_df = train_test_split(df, test_size=5, random_state=42)
+
+        os.makedirs("outputs", exist_ok=True)
+
+        # 6) Generate predictions on train split → outputs/inference.tsv
+        print("Generating predictions for training set...")
+        train_preds = []
+        for idx, doctor_note in enumerate(train_df["doctor_note"]):
+            print(f"Predicting train {idx+1}/{len(train_df)}...")
+            pred_soap = convert_to_soap(doctor_note)
+            train_preds.append(pred_soap)
+            
+            if (idx + 1) % 3 == 0:
+                torch.cuda.empty_cache()
+                gc.collect()
+
+        inf_df = train_df.reset_index(drop=True).copy()
+        inf_df["id"] = inf_df.index + 1
+        inf_df["ground_truth_soap"] = inf_df["soap_note"]
+        inf_df["predicted_soap"] = train_preds
         
-        # Clean up model after use
+        # Save inference results
+        inf_df[["id", "ground_truth_soap", "predicted_soap"]].to_csv(
+            "outputs/inference.tsv", sep="\t", index=False
+        )
+        print("Inference results saved!")
+
+        # 7) Generate predictions on test split → outputs/eval.csv
+        print("Generating predictions for test set...")
+        test_preds = []
+        for idx, doctor_note in enumerate(test_df["doctor_note"]):
+            print(f"Predicting test {idx+1}/{len(test_df)}...")
+            pred_soap = convert_to_soap(doctor_note)
+            test_preds.append(pred_soap)
+            
+            torch.cuda.empty_cache()
+            gc.collect()
+
+        eval_df = pd.DataFrame({
+            "id": range(1, len(test_preds) + 1),
+            "predicted_soap": test_preds
+        })
+        eval_df.to_csv("outputs/eval.csv", index=False)
+        print("Evaluation results saved!")
+
+        # 8) Save complete dataset for reference
+        complete_df = pd.DataFrame({
+            "id": range(1, len(docs) + 1),
+            "doctor_note": docs,
+            "soap_note": soaps
+        })
+        complete_df.to_csv("outputs/complete_dataset.csv", index=False)
+        print("Complete dataset saved!")
+
+        # 9) Cleanup model
         del model
         torch.cuda.empty_cache()
         gc.collect()
-        
+        print("Model cleaned up!")
 
+        # 10) Return status + file paths for download
+        return (
+            f"✅ Successfully generated 20 notes!\n"
+            f"📊 Training set: {len(train_df)} notes\n"
+            f"🧪 Test set: {len(test_df)} notes\n"
+            f"💾 Files ready for download",
+            "outputs/inference.tsv",
+            "outputs/eval.csv",
+            "outputs/complete_dataset.csv"
+        )
 
     except Exception as e:
         traceback.print_exc()
-        return (f"❌ Error: {e}", None, None)
+        return (f"❌ Error: {e}", None, None, None)
 
 # ─── Gradio UI ─────────────────────────────────────────────────────────────────
-with gr.Blocks() as demo:
-    gr.Markdown("# Gemma‑3n SOAP Generator 🩺")
-    btn      = gr.Button("Generate & Export 20 Notes")
-    status   = gr.Textbox(interactive=False, label="Status")
-    inf_file = gr.File(label="Download inference.tsv")
-    eval_file= gr.File(label="Download eval.csv")
+with gr.Blocks(title="SOAP Generator") as demo:
+    gr.Markdown("""
+    # 🩺 Medical SOAP Note Generator
+    
+    This app generates realistic doctor's notes and converts them to SOAP format:
+    - **S**ubjective: Patient's reported symptoms
+    - **O**bjective: Observable findings and exam results  
+    - **A**ssessment: Clinical diagnosis/impression
+    - **P**lan: Treatment plan and follow-up
+    
+    **Process:**
+    1. Generate 20 realistic doctor's progress notes
+    2. Convert each note to structured SOAP format
+    3. Split into 15 training + 5 test samples
+    4. Generate predictions and export files
+    """)
+    
+    with gr.Row():
+        btn = gr.Button("🚀 Generate & Export SOAP Notes", variant="primary", size="lg")
+    
+    with gr.Row():
+        status = gr.Textbox(
+            label="📋 Generation Status", 
+            interactive=False, 
+            lines=5,
+            placeholder="Click the button above to start generation..."
+        )
+    
+    with gr.Row():
+        with gr.Column():
+            inf_file = gr.File(label="📊 Download inference.tsv (Training Predictions)")
+        with gr.Column():
+            eval_file = gr.File(label="🧪 Download eval.csv (Test Predictions)")
+        with gr.Column():
+            complete_file = gr.File(label="💾 Download complete_dataset.csv (All Data)")
 
     btn.click(
         fn=generate_and_export,
         inputs=None,
-        outputs=[status, inf_file, eval_file]
+        outputs=[status, inf_file, eval_file, complete_file]
     )
 
+    gr.Markdown("""
+    ### 📁 Output Files:
+    - **inference.tsv**: Training set with ground truth and predicted SOAP notes
+    - **eval.csv**: Test set predictions only
+    - **complete_dataset.csv**: All 20 generated doctor notes and SOAP conversions
+    """)
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()