Update app.py

app.py

@@ -1,140 +1,93 @@
 import os
-import torch
 import pandas as pd
-import easyocr
 import gradio as gr
-from transformers import (
-    AutoConfig,
-    AutoProcessor,
-    AutoTokenizer,
-    AutoModelForImageTextToText
-)
+from kaggle_secrets import UserSecretsClient
+from transformers import AutoProcessor, AutoTokenizer, AutoModelForImageTextToText
+from sklearn.model_selection import train_test_split
+import torch
 
-MODEL_ID = "google/gemma-3n-e2b-it"
-HF_TOKEN = os.environ.get("HF_TOKEN")  # set via Space secrets
+HF_TOKEN = UserSecretsClient().get_secret("HF_TOKEN")
+MODEL_ID  = "google/gemma-3n-e2b-it"
 
-# 1) Load the model and OCR reader
-config    = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True, use_auth_token=HF_TOKEN)
+# 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     = AutoModelForImageTextToText.from_pretrained(
-    MODEL_ID,
-    config=config,
-    trust_remote_code=True,
-    use_auth_token=HF_TOKEN,
-    load_in_8bit=True,
-    device_map="auto"
-)
-device = next(model.parameters()).device
-
-ocr_reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available(), verbose=False)
-
-def generate_soap_note(text: str) -> str:
-    prompt = f"""You are a medical AI assistant. Convert these notes into a SOAP note:
-
-{text}
-
-Format as:
-S - SUBJECTIVE:
-O - OBJECTIVE:
-A - ASSESSMENT:
-P - PLAN:
 
-Produce the complete SOAP."""
-    inputs = processor.apply_chat_template(
-        [
-            {"role":"system","content":[{"type":"text","text":"Expert medical AI assistant."}]},
-            {"role":"user",  "content":[{"type":"text","text":prompt}]}
-        ],
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt"
-    ).to(device)
-    input_len = inputs["input_ids"].shape[-1]
-    with torch.no_grad():
+def generate_all_notes():
+    # 1) Load the full 8‑bit model on demand
+    model = AutoModelForImageTextToText.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True,
+        use_auth_token=HF_TOKEN,
+        load_in_8bit=True,
+        device_map="auto"
+    )
+    device = next(model.parameters()).device
+
+    # helper to turn text→SOAP
+    def to_soap(text):
+        inputs = processor.apply_chat_template(
+            [
+                {"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,
+            return_tensors="pt",
+            return_dict=True
+        ).to(device)
         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,
-            disable_compile=True
+            top_p=0.95,
+            pad_token_id=processor.tokenizer.eos_token_id
         )
-    return processor.batch_decode(out[:, input_len:], skip_special_tokens=True)[0].strip()
-
-# 2) On startup: generate 100 synthetic note pairs and save TSVs
-docs, soaps = [], []
-for i in range(1, 101):
-    doc = generate_soap_note("Generate a realistic, concise doctor's progress note for a single patient encounter.")
-    docs.append(doc)
-    soaps.append(generate_soap_note(doc))
-    if i % 10 == 0:
-        print(f"Generated {i}/100")
-
-pd.DataFrame({"doc_note": docs}).to_csv("doc_notes.tsv", sep="\t", index=False)
-pd.DataFrame({"soap_note": soaps}).to_csv("ground_truth_soap.tsv", sep="\t", index=False)
-print("✅ Saved doc_notes.tsv & ground_truth_soap.tsv")
-
-# Colab/Kaggle/Space Cell: Split, infer on train (70), infer on test (30), and save files
-
-import os
-import pandas as pd
-from sklearn.model_selection import train_test_split
-
-# Make outputs directory
-os.makedirs("outputs", exist_ok=True)
-
-# 1) Load generated notes and ground truths
-docs = pd.read_csv("doc_notes.tsv", sep="\t")            # has column 'doc_note'
-gts  = pd.read_csv("ground_truth_soap.tsv", sep="\t")    # has column 'soap_note'
-
-full = pd.DataFrame({
-    "doc_note":           docs["doc_note"],
-    "ground_truth_soap":  gts["soap_note"]
-})
-
-# 2) Split 70% train / 30% test
-train_df, test_df = train_test_split(full, test_size=0.3, random_state=42)
-
-# 3) Run inference on the 70‑row train split and save inference.tsv
-train_preds = []
-for idx, row in train_df.reset_index(drop=True).iterrows():
-    pred = generate_soap_note(row["doc_note"])
-    train_preds.append(pred)
-
-inference_df = train_df.reset_index(drop=True).copy()
-inference_df["id"]             = inference_df.index + 1
-inference_df["predicted_soap"] = train_preds
-inference_df = inference_df[["id","ground_truth_soap","predicted_soap"]]
-inference_df.to_csv("outputs/inference.tsv", sep="\t", index=False)
-
-# 4) Run inference on the 30��row test split and save eval.csv
-test_preds = []
-for idx, row in test_df.reset_index(drop=True).iterrows():
-    pred = generate_soap_note(row["doc_note"])
-    test_preds.append(pred)
-
-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("✅ Saved:")
-print("  outputs/inference.tsv (70 rows: id, ground_truth_soap, predicted_soap)")
-print("  outputs/eval.csv      (30 rows: id, predicted_soap)")
-
-
-# 3) Blank Gradio UI placeholder
-def noop(): 
-    return "Data generated — check TSV files in the repo."
+        return processor.batch_decode(out[:, inputs["input_ids"].shape[-1]:], skip_special_tokens=True)[0].strip()
+
+    # 2) Generate 100 raw docs + ground truths
+    docs, gts = [], []
+    for i in range(100):
+        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:
+            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)
+
+    os.makedirs("outputs", exist_ok=True)
+
+    # 4) Inference on train split → 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
+    )
+
+    # 5) Inference on test split → 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 (
+        "✅ Done!\n"
+        f"– outputs/inference.tsv (70 rows with id, GT & pred)\n"
+        f"– outputs/eval.csv      (30 rows with id & pred)"
+    )
 
 with gr.Blocks() as demo:
-    gr.Markdown("# SOAP Generator Space")
-    gr.Button("Generate Completed (see logs)").click(noop, [], "output")
-    gr.Textbox(label="Status", interactive=False, lines=2, placeholder="Ready", elem_id="output")
+    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)
 
-if __name__ == "__main__":
-    demo.queue().launch(server_name="0.0.0.0", server_port=7860)
+if __name__=="__main__":
+    demo.launch()