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

model_id = "Yui-Father/yui-lora-model"
dtype = torch.float16

try:
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype=dtype,
        load_in_8bit=True,
        device_map="auto",
        trust_remote_code=True,
    )
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {e}")
    raise e

def generate_response(message, history):
    formatted_prompt = f"أنت: {message}\nيوي:"
    inputs = tokenizer(formatted_prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        inputs.input_ids,
        max_new_tokens=200,
        do_sample=True,
        top_p=0.7,
        temperature=0.7,
        num_beams=1,
        pad_token_id=tokenizer.eos_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)

    if "أنت:" in response:
        response = response.split("أنت:")[0].strip()

    return response

iface = gr.ChatInterface(
    generate_response,
    title="تحدث مع يوي (نموذجك الخاص)",
    description="مرحباً بك! أنا يوي، نموذجك الذكي الخاص، جاهزة للمحادثة.",
    examples=[["كيف حالك؟"], ["ما هو طعامك المفضل؟"], ["احكي لي قصة قصيرة."]],
)

iface.launch()