Update app.py

app.py

@@ -13,7 +13,7 @@ from transformers import (
 )
 from sklearn.model_selection import train_test_split
 
-# ─── Silence unrecognized‐flag warnings ────────────────────────────────────────
+# ─── Silence irrelevant warnings ───────────────────────────────────────────────
 logging.set_verbosity_error()
 
 # ─── Configuration ────────────────────────────────────────────────────────────
@@ -22,7 +22,7 @@ if not HF_TOKEN:
     raise RuntimeError("Missing HF_TOKEN in env vars – set it under Space Settings → Secrets")
 MODEL_ID = "google/gemma-3n-e2b-it"
 
-# ─── Fast startup: load only the small pieces ──────────────────────────────────
+# ─── Fast startup: load only processor & tokenizer ─────────────────────────────
 processor = AutoProcessor.from_pretrained(
     MODEL_ID, trust_remote_code=True, token=HF_TOKEN
 )
@@ -30,10 +30,10 @@ tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID, trust_remote_code=True, token=HF_TOKEN
 )
 
-# ─── Heavy work deferred until button click ───────────────────────────────────
+# ─── Heavy work runs on button click ───────────────────────────────────────────
 def generate_and_export():
     try:
-        # 1) Lazy-load the full FP16 model (heavy)
+        # 1) Lazy‑load the full FP16 model
         model = AutoModelForImageTextToText.from_pretrained(
             MODEL_ID,
             trust_remote_code=True,
@@ -43,7 +43,7 @@ def generate_and_export():
         )
         device = next(model.parameters()).device
 
-        # 2) Helper to generate a SOAP note from arbitrary text
+        # 2) Text→SOAP helper
         def to_soap(text: str) -> str:
             inputs = processor.apply_chat_template(
                 [
@@ -62,30 +62,29 @@ def generate_and_export():
                 top_p=0.95,
                 temperature=0.1,
                 pad_token_id=processor.tokenizer.eos_token_id,
-                use_cache=False      # disable HybridCache
+                use_cache=False
             )
             prompt_len = inputs["input_ids"].shape[-1]
             return processor.batch_decode(
                 out[:, prompt_len:], skip_special_tokens=True
             )[0].strip()
 
-        # 3) Generate 100 synthetic doc notes & ground-truth SOAPs
+        # 3) Generate 20 doc notes + ground truths
         docs, gts = [], []
-        for i in range(1, 101):
+        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 % 20 == 0:
+            if i % 5 == 0:
                 torch.cuda.empty_cache()
 
-        # 4) 70/30 train-test split
+        # 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=0.3, random_state=42)
+        train_df, test_df = train_test_split(df, test_size=5, random_state=42)
 
-        # ensure outputs folder
         os.makedirs("outputs", exist_ok=True)
 
-        # 5) Run inference on train split → inference.tsv
+        # 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
@@ -94,30 +93,28 @@ def generate_and_export():
             "outputs/inference.tsv", sep="\t", index=False
         )
 
-        # 6) Run inference on test split → eval.csv
+        # 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),
+            "id":             range(1, len(test_preds) + 1),
             "predicted_soap": test_preds
         }).to_csv("outputs/eval.csv", index=False)
 
-        # 7) Success: return status + file paths
+        # 7) Return status + file paths for download
         return (
-            "✅ Generation complete! Download below ⤵",
+            "✅ Done with 20 notes (15 train / 5 test)!",
             "outputs/inference.tsv",
             "outputs/eval.csv"
         )
 
     except Exception as e:
-        # Print full traceback to the Space logs
         traceback.print_exc()
-        # Return the error message to the UI
         return (f"❌ Error: {e}", None, None)
 
 # ─── Gradio UI ──────────────────��──────────────────────────────────────────────
 with gr.Blocks() as demo:
     gr.Markdown("# Gemma‑3n SOAP Generator 🩺")
-    btn      = gr.Button("Generate & Export 100 Notes")
+    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")