Update app.py

app.py

@@ -2,23 +2,19 @@
 
 import os
 import pandas as pd
-import gradio as gr
 import torch
-from transformers import (
-    AutoProcessor,
-    AutoTokenizer,
-    AutoModelForImageTextToText
-)
+import gradio as gr
+from transformers import AutoProcessor, AutoTokenizer, AutoModelForImageTextToText
 from sklearn.model_selection import train_test_split
 
-# 1) Retrieve your HF_TOKEN from environment (set in Space Settings → Secrets)
+# 1) Configuration
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if not HF_TOKEN:
-    raise RuntimeError("Missing HF_TOKEN env var! Please add it in your Space settings → Secrets.")
+    raise RuntimeError("Missing HF_TOKEN in env vars – please add it under Settings → Secrets")
 
 MODEL_ID = "google/gemma-3n-e2b-it"
 
-# 2) Eagerly load the small bits (processor & tokenizer) so the UI starts fast
+# 2) Eagerly load only the processor & tokenizer (fast startup)
 processor = AutoProcessor.from_pretrained(
     MODEL_ID, trust_remote_code=True, token=HF_TOKEN
 )
@@ -26,9 +22,12 @@ tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID, trust_remote_code=True, token=HF_TOKEN
 )
 
-def generate_all_and_split():
-    """Called when the user clicks the button—loads full model, generates & saves files."""
-    # a) Lazy‑load the 8‑bit quantized model (heavy)
+def generate_and_export():
+    """
+    On button click: lazily load the 8‑bit model, generate 100 doc→SOAP pairs,
+    split 70/30, run inference & eval, write files, and return download links.
+    """
+    # a) Load full model in 8‑bit
     model = AutoModelForImageTextToText.from_pretrained(
         MODEL_ID,
         trust_remote_code=True,
@@ -39,6 +38,7 @@ def generate_all_and_split():
     device = next(model.parameters()).device
 
     def to_soap(text: str) -> str:
+        # wrap the chat‐template + generate call
         inputs = processor.apply_chat_template(
             [
                 {"role":"system","content":[{"type":"text","text":"You are a medical AI assistant."}]},
@@ -57,10 +57,11 @@ def generate_all_and_split():
             temperature=0.1,
             pad_token_id=processor.tokenizer.eos_token_id
         )
+        # strip off prompt tokens
         prompt_len = inputs["input_ids"].shape[-1]
         return processor.batch_decode(out[:, prompt_len:], skip_special_tokens=True)[0].strip()
 
-    # b) Generate 100 doc_notes + ground_truth SOAPs
+    # b) Generate 100 synthetic doc_notes & ground_truth SOAPs
     docs, gts = [], []
     for i in range(1, 101):
         doc = to_soap("Generate a realistic, concise doctor's progress note for a single patient encounter.")
@@ -69,13 +70,14 @@ def generate_all_and_split():
         if i % 20 == 0:
             torch.cuda.empty_cache()
 
-    # c) Split 70/30
+    # c) 70/30 split
     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)
 
+    # ensure outputs dir exists
     os.makedirs("outputs", exist_ok=True)
 
-    # d) Inference on train → inference.tsv
+    # d) 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
@@ -84,25 +86,33 @@ def generate_all_and_split():
         "outputs/inference.tsv", sep="\t", index=False
     )
 
-    # e) Inference on test → eval.csv
+    # e) 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)
 
+    # return status + file paths for download
     return (
-        "✅ Done!\n"
-        f"• outputs/inference.tsv (70 rows with id, GT, pred)\n"
-        f"• outputs/eval.csv      (30 rows with id, pred)"
+        "✅ Generation complete!",
+        "outputs/inference.tsv",
+        "outputs/eval.csv"
     )
 
-# 3) Gradio UI—instant startup
+# 3) Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("## Gemma‑3n SOAP Generator 🩺")
-    btn = gr.Button("Generate & Save 100 Notes → 70/30 Split → inference & eval")
-    status = gr.Textbox(interactive=False, label="Status")
-    btn.click(fn=generate_all_and_split, inputs=None, outputs=status)
+    gr.Markdown("# Gemma‑3n SOAP Generator 🩺")
+    generate_btn = gr.Button("Generate & Export 100 Notes")
+    status      = gr.Textbox(interactive=False, label="Status")
+    inf_file    = gr.File(label="Download inference.tsv")
+    eval_file   = gr.File(label="Download eval.csv")
+
+    generate_btn.click(
+        fn=generate_and_export,
+        inputs=None,
+        outputs=[status, inf_file, eval_file]
+    )
 
 if __name__ == "__main__":
     demo.launch()