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

# 🔥 تحديد الجهاز المناسب
device = "cuda" if torch.cuda.is_available() else "cpu"

# 🔥 تحميل الموديل بطريقة محسنة
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,  # ✅ ضبط الـ device_map بشكل أدق
    trust_remote_code=True
).eval()

# 🔥 تحميل التوكنيزر
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token  

# 🔥 إعداد الدالة الخاصة بالـ Chatbot
def chatbot(user_input):
    if not user_input.strip():
        return "Please enter a message."

    inputs = tokenizer(user_input, return_tensors="pt", padding=True).to(device)

    with torch.no_grad():
        output = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],  # ✅ إصلاح مشكلة الماسك
            max_new_tokens=50,  # ✅ استخدام max_new_tokens بدلاً من max_length
            temperature=0.6,
            top_p=0.8,
            do_sample=True,
            early_stopping=True,
            pad_token_id=tokenizer.eos_token_id
        )

    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# 🔥 تشغيل الواجهة بدون share=True
iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="Octagon 2.0 Chatbot")

# ✅ إصلاح مشكلة POST method not allowed
iface.launch(ssl=False, debug=True)