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()