Update app.py

app.py

@@ -1,31 +1,44 @@
+# app.py
+
 import os
 import pandas as pd
 import gradio as gr
-from kaggle_secrets import UserSecretsClient
-from transformers import AutoProcessor, AutoTokenizer, AutoModelForImageTextToText
-from sklearn.model_selection import train_test_split
 import torch
+from transformers import (
+    AutoProcessor,
+    AutoTokenizer,
+    AutoModelForImageTextToText
+)
+from sklearn.model_selection import train_test_split
 
-HF_TOKEN = UserSecretsClient().get_secret("HF_TOKEN")
-MODEL_ID  = "google/gemma-3n-e2b-it"
+# 1) Retrieve your HF_TOKEN from environment (set in Space Settings → Secrets)
+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.")
 
-# Only load small pieces at startup
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, use_auth_token=HF_TOKEN)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, use_auth_token=HF_TOKEN)
+MODEL_ID = "google/gemma-3n-e2b-it"
 
-def generate_all_notes():
-    # 1) Load the full 8‑bit model on demand
+# 2) Eagerly load the small bits (processor & tokenizer) so the UI starts fast
+processor = AutoProcessor.from_pretrained(
+    MODEL_ID, trust_remote_code=True, token=HF_TOKEN
+)
+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)
     model = AutoModelForImageTextToText.from_pretrained(
         MODEL_ID,
         trust_remote_code=True,
-        use_auth_token=HF_TOKEN,
+        token=HF_TOKEN,
         load_in_8bit=True,
         device_map="auto"
     )
     device = next(model.parameters()).device
 
-    # helper to turn text→SOAP
-    def to_soap(text):
+    def to_soap(text: str) -> str:
         inputs = processor.apply_chat_template(
             [
                 {"role":"system","content":[{"type":"text","text":"You are a medical AI assistant."}]},
@@ -40,28 +53,29 @@ def generate_all_notes():
             **inputs,
             max_new_tokens=400,
             do_sample=True,
-            temperature=0.1,
             top_p=0.95,
+            temperature=0.1,
             pad_token_id=processor.tokenizer.eos_token_id
         )
-        return processor.batch_decode(out[:, inputs["input_ids"].shape[-1]:], skip_special_tokens=True)[0].strip()
+        prompt_len = inputs["input_ids"].shape[-1]
+        return processor.batch_decode(out[:, prompt_len:], skip_special_tokens=True)[0].strip()
 
-    # 2) Generate 100 raw docs + ground truths
+    # b) Generate 100 doc_notes + ground_truth SOAPs
     docs, gts = [], []
-    for i in range(100):
+    for i in range(1, 101):
         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+1) % 20 == 0:
+        if i % 20 == 0:
             torch.cuda.empty_cache()
 
-    # 3) Split 70/30
-    full_df = pd.DataFrame({"doc_note": docs, "ground_truth_soap": gts})
-    train_df, test_df = train_test_split(full_df, test_size=0.3, random_state=42)
+    # c) 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)
 
     os.makedirs("outputs", exist_ok=True)
 
-    # 4) Inference on train split → inference.tsv
+    # d) 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
@@ -70,24 +84,25 @@ def generate_all_notes():
         "outputs/inference.tsv", sep="\t", index=False
     )
 
-    # 5) Inference on test split → eval.csv
+    # e) 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 (
         "✅ Done!\n"
-        f"– outputs/inference.tsv (70 rows with id, GT & pred)\n"
-        f"– outputs/eval.csv      (30 rows with id & pred)"
+        f"• outputs/inference.tsv (70 rows with id, GT, pred)\n"
+        f"• outputs/eval.csv      (30 rows with id, pred)"
     )
 
+# 3) Gradio UI—instant startup
 with gr.Blocks() as demo:
-    gr.Markdown("## Gemma‑3n SOAP Generator")
-    btn = gr.Button("Generate 100 → split 70/30 → inference & eval")
-    out = gr.Textbox(interactive=False, label="Status")
-    btn.click(fn=generate_all_notes, inputs=None, outputs=out)
+    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)
 
-if __name__=="__main__":
+if __name__ == "__main__":
     demo.launch()