# main.py

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

app = FastAPI()

# تعريف نموذج البيانات للطلب: يحتوي على حقل واحد "input" (النص المرسل من العميل)
class GenerateRequest(BaseModel):
    input: str  # جملة المُدخل (مثال: "اسمي محمد")

# تحميل النموذج والـ tokenizer من مكتبة Hugging Face
# هنا نستخدم مثال akhooli/gpt2-small-arabic (137M معلمات):contentReference[oaicite:3]{index=3}.
# نستخدم device=-1 لاستخدام الـ CPU فقط (Hot GPU غير متاح في خطة Spaces المجانية).
model_name = "akhooli/gpt2-small-arabic"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)

@app.post("/generate")
def generate_greeting(request: GenerateRequest):
    # استخراج النص المرسل في الطلب
    user_input = request.input.strip()
    if not user_input:
        # في حال كان الحقل فارغاً، نرجع رسالة خطأ بسيطة
        return {"response": "الرجاء تزويد جملة صحيحة في الحقل 'input'."}

    # استخدم النموذج لتوليد نص ترحيبي بناءً على المدخل
    # نحدد بعض معايير التوليد مثل max_new_tokens و top_p لتحسين الجودة.
    result = generator(
        user_input,
        max_new_tokens=50,      # الحد الأقصى لعدد الرموز الجديدة المولدة
        do_sample=True,         # استخدام التوليد العشوائي (Sampling) بدلًا من التوليد الحتمي
        top_p=0.95,             # لتحسين التنوع
        return_full_text=False  # لإرجاع الجزء المولد فقط (بدون نص المدخل الأصلي)
    )

    # الحصول على النص المولد وترتيبه
    generated_text = result[0]['generated_text'].strip()

    # في بعض الأحيان قد يبدأ النص المولد بحرف بليغ أو فاصلة، لذا نزيل أي علامات ترقيم غير ضرورية في البداية
    greeting = generated_text
    # تأكد من أن الرد ليس فارغاً تماماً
    if greeting == "":
        greeting = "مرحباً!"

    return {"response": greeting}