Update app.py

app.py

@@ -1,20 +1,23 @@
-# app.py
-
 import os
 import pandas as pd
 import torch
 import gradio as gr
-from transformers import AutoProcessor, AutoTokenizer, AutoModelForImageTextToText
+from transformers import (
+    AutoProcessor,
+    AutoTokenizer,
+    AutoModelForImageTextToText
+)
 from sklearn.model_selection import train_test_split
 
-# 1) Configuration
 HF_TOKEN = os.environ.get("HF_TOKEN")
 if not HF_TOKEN:
-    raise RuntimeError("Missing HF_TOKEN in env vars – please add it under Settings → Secrets")
+    raise RuntimeError(
+        "Missing HF_TOKEN in env vars – add it under Settings → Secrets"
+    )
 
 MODEL_ID = "google/gemma-3n-e2b-it"
 
-# 2) Eagerly load only the processor & tokenizer (fast startup)
+# Load processor & tokenizer at top level for fast startup
 processor = AutoProcessor.from_pretrained(
     MODEL_ID, trust_remote_code=True, token=HF_TOKEN
 )
@@ -23,25 +26,25 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 
 def generate_and_export():
-    out = model.generate(
-    **inputs,
-    max_new_tokens=400,
-    do_sample=True,
-    top_p=0.95,
-    temperature=0.1,
-    pad_token_id=processor.tokenizer.eos_token_id,
-    use_cache=False         # ← disable HybridCache
-)
-
-
+    """
+    On button click: load full model, generate 100 notes,
+    split 70/30, run inference & eval, save files, return download links.
+    """
+    # Load the heavy model here
+    model = AutoModelForImageTextToText.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True,
+        token=HF_TOKEN,
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
     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."}]},
-                {"role":"user",  "content":[{"type":"text","text":text}]}
+                {"role": "system", "content":[{"type":"text","text":"You are a medical AI assistant."}]},
+                {"role": "user",   "content":[{"type":"text","text":text}]}
             ],
             add_generation_prompt=True,
             tokenize=True,
@@ -54,60 +57,63 @@ def generate_and_export():
             do_sample=True,
             top_p=0.95,
             temperature=0.1,
-            pad_token_id=processor.tokenizer.eos_token_id
+            pad_token_id=processor.tokenizer.eos_token_id,
+            use_cache=False
         )
-        # strip off prompt tokens
         prompt_len = inputs["input_ids"].shape[-1]
-        return processor.batch_decode(out[:, prompt_len:], skip_special_tokens=True)[0].strip()
+        return processor.batch_decode(
+            out[:, prompt_len:], skip_special_tokens=True
+        )[0].strip()
 
-    # b) Generate 100 synthetic doc_notes & ground_truth SOAPs
+    # Generate 100 docs + GTs
     docs, gts = [], []
     for i in range(1, 101):
-        doc = to_soap("Generate a realistic, concise doctor's progress note for a single patient encounter.")
+        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:
             torch.cuda.empty_cache()
 
-    # c) 70/30 split
+    # Split 70/30
     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
+    # Ensure outputs dir
     os.makedirs("outputs", exist_ok=True)
 
-    # d) Inference on train split → outputs/inference.tsv
+    # Inference on train → 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["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
     )
 
-    # e) Inference on test split → outputs/eval.csv
+    # Inference on test → 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)
 
-    # return status + file paths for download
     return (
-        "✅ Generation complete!",
+        "✅ Done!",
         "outputs/inference.tsv",
         "outputs/eval.csv"
     )
 
-# 3) Gradio UI
+# Build Gradio interface (starts immediately)
 with gr.Blocks() as demo:
     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")
+    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(
+    btn.click(
         fn=generate_and_export,
         inputs=None,
         outputs=[status, inf_file, eval_file]