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

# تحميل الموديل والتوكنيزر
model_name = "NousResearch/Hermes-2-Pro-Mistral-7B"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    device_map="auto"
)

# تجهيز الـ pipeline
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# دالة الرد على المستخدم
def chatbot(user_input):
    prompt = f"""### النظام:
أنت مساعد ذكي متخصص في مجال طب الأسنان، تجيب على الأسئلة الطبية بدقة وباللغة العربية إن أمكن.
### المستخدم:
{user_input}
### المساعد:"""
    
    response = generator(
        prompt,
        max_new_tokens=300,
        do_sample=True,
        temperature=0.7,
        top_p=0.95
    )
    
    output_text = response[0]["generated_text"]
    # فصل إجابة المساعد من النص الكامل
    answer = output_text.split("### المساعد:")[-1].strip()
    return answer

# تصميم الواجهة بـ Gradio
gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(lines=3, label="اسأل سؤالاً في طب الأسنان"),
    outputs=gr.Textbox(lines=8, label="إجابة المساعد"),
    title="طبيب الأسنان الآلي",
    description="مساعد ذكي يجيب على الأسئلة الطبية المتخصصة في مجال طب الأسنان"
).launch()