import os
import zipfile
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
from peft import PeftModel
from huggingface_hub import login

# --- گام ۱: احراز هویت Hugging Face
hf_token = os.environ.get("HF_TOKEN")
if not hf_token:
    raise ValueError("❌ HF_TOKEN not found in environment secrets.")
login(hf_token)

# اگر فایل اشتباه وجود داره و فایل جدید نه
if os.path.exists("dorna-diabetes-finetuned-20250514T183411Z-1-001.zip") and not os.path.exists("dorna-diabetes-finetuned.zip"):
    os.rename("dorna-diabetes-finetuned-20250514T183411Z-1-001.zip", "dorna-diabetes-finetuned.zip")
    print("✅ اسم فایل تغییر کرد.")

# --- گام ۲: اکسترکت فایل فشرده (فقط بار اول)
if not os.path.exists("dorna-diabetes-finetuned"):
    with zipfile.ZipFile("dorna-diabetes-finetuned.zip", "r") as zip_ref:
        zip_ref.extractall(".")
    print("✅ فایل LoRA اکسترکت شد.")

# --- گام ۳: بارگذاری مدل پایه و LoRA
BASE_MODEL = "PartAI/Dorna-Llama3-8B-Instruct"
LORA_PATH = "./dorna-diabetes-finetuned"

print("🔹 در حال بارگذاری مدل پایه...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_auth_token=hf_token)
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    token=HF_TOKEN,
    trust_remote_code=True
)

print("🔹 در حال بارگذاری LoRA...")
model = PeftModel.from_pretrained(base_model, LORA_PATH)
model.eval()

streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

# --- گام ۴: رابط چت با Gradio
def chat(prompt):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
    with torch.no_grad():
        output = model.generate(
            input_ids=input_ids,
            max_new_tokens=200,
            temperature=0.7,
            top_p=0.9,
            do_sample=True
        )
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response[len(prompt):].strip()

iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="💬 Dorna LoRA Model")

iface.launch()