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

# Load base model dan tokenizer
base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
lora_model = "Wiefdw/modelAnevia-TinyLlama-LoRA-v2"  # Ganti dengan repo kamu

tokenizer = AutoTokenizer.from_pretrained(lora_model)
base = AutoModelForCausalLM.from_pretrained(
    base_model,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

# Load LoRA adapter
model = PeftModel.from_pretrained(base, lora_model)
model.eval()

# Fungsi chat
def chat_with_model(prompt, max_tokens=200, temperature=0.7):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=True,
            top_p=0.95,
            eos_token_id=tokenizer.eos_token_id
        )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Gradio Interface
demo = gr.Interface(
    fn=chat_with_model,
    inputs=[
        gr.Textbox(label="Masukkan pertanyaan Anda", lines=3, placeholder="Contoh: Saya sering lemas dan pusing..."),
        gr.Slider(50, 500, value=200, label="Max Tokens"),
        gr.Slider(0.1, 1.0, value=0.7, label="Temperature")
    ],
    outputs=gr.Textbox(label="Respon Model"),
    title="💉 Anemia Chatbot (TinyLlama + LoRA)",
    description="Model TinyLlama yang di-fine-tune dengan LoRA untuk percakapan seputar anemia dan gejala kesehatannya."
)

demo.launch()