import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "ViteMamba/TigerLLM-Medical-Bengali"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True,
    device_map="cpu"
)
model.eval()
print("✅ Ready!")

def ask(question, max_tokens=256, temperature=0.7):
    if not question.strip():
        return "অনুগ্রহ করে একটি প্রশ্ন লিখুন।"
    prompt = f"""<bos><start_of_turn>system
আপনি একজন বাংলা চিকিৎসা সহকারী। নিচের প্রশ্নের বিস্তারিত উত্তর দিন।<end_of_turn>
<start_of_turn>user
{question}<end_of_turn>
<start_of_turn>model
"""
    inputs = tokenizer(prompt, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            temperature=float(temperature),
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    generated = outputs[0][inputs["input_ids"].shape[1]:]
    return tokenizer.decode(generated, skip_special_tokens=True)

demo = gr.Interface(
    fn=ask,
    inputs=[
        gr.Textbox(
            label="বাংলা চিকিৎসা প্রশ্ন",
            placeholder="যেমন: ডায়াবেটিসের লক্ষণ কী?",
            lines=3
        ),
        gr.Slider(50, 512, value=256, step=1, label="Max Tokens"),
        gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature"),
    ],
    outputs=gr.Textbox(label="উত্তর", lines=10),
    title="🏥 TigerLLM Bengali Medical Assistant",
    description="বাংলায় চিকিৎসা প্রশ্নের উত্তর পান।",
    examples=[
        ["ডায়াবেটিসের প্রধান লক্ষণগুলো কী কী?"],
        ["উচ্চ রক্তচাপ কমানোর উপায় কী?"],
        ["নিউমোনিয়া কীভাবে ছড়ায়?"],
        ["হার্ট অ্যাটাকের লক্ষণগুলো কী?"],
    ],
)

demo.launch(server_name="0.0.0.0", server_port=7860)