Create app.py

transformers
peft
bitsandbytes
accelerate
torch
gradio
spaces

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from peft import PeftModel
+import os
+
+BASE_MODEL = "Qwen/Qwen2.5-Math-72B-Instruct"
+ADAPTER_REPO = "billwang37/mathbio-qwen-72b-round2"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Load in 4-bit to fit on ZeroGPU H200
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_quant_type="nf4",
+)
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO, token=HF_TOKEN)
+
+print("Loading base model in 4-bit...")
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    quantization_config=bnb_config,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+print("Loading LoRA adapter...")
+model = PeftModel.from_pretrained(base_model, ADAPTER_REPO, token=HF_TOKEN)
+model.eval()
+print("Model ready.")
+
+SYSTEM_PROMPT = "You are MathBioAgent, an expert AI assistant specialized in mathematical biology, epidemiology, operator learning, and partial differential equations."
+
+@spaces.GPU(duration=120)
+def chat(message, history):
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for h in history:
+        messages.append({"role": "user", "content": h[0]})
+        messages.append({"role": "assistant", "content": h[1]})
+    messages.append({"role": "user", "content": message})
+    
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=1024,
+            temperature=0.3,
+            do_sample=True,
+            top_p=0.9,
+        )
+    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    return response
+
+demo = gr.ChatInterface(
+    fn=chat,
+    title="MathBio AI — Specialized for Mathematical Biology",
+    description="Ask questions about epidemic modeling, PDEs, operator learning, and mathematical biology. Based on Qwen2.5-Math-72B fine-tuned on 27K research examples.",
+    examples=[
+        "What is R0 for an SIR model with beta=0.4 and gamma=0.1?",
+        "Derive the stability condition for the SEIR endemic equilibrium.",
+        "Explain the Keller-Segel chemotaxis model.",
+    ],
+)
+
+demo.launch()