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

BASE_MODEL = "PartAI/Dorna-Llama3-8B-Instruct"
LORA_PATH = "./dorna-diabetes-finetuned"  # این پوشه رو آپلود می‌کنی توی اسپیس

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    torch_dtype=torch.float16,
)

model = PeftModel.from_pretrained(base_model, LORA_PATH)

def generate_response(prompt):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    with torch.no_grad():
        output = model.generate(
            input_ids=input_ids,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
        )
    return tokenizer.decode(output[0], skip_special_tokens=True)

gr.Interface(fn=generate_response, inputs="text", outputs="text", title="Dorna-Llama3 LoRA").launch()