from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import gradio as gr

# 1. تهيئة تطبيق FastAPI
app = FastAPI()

# 2. تحميل النموذج مرة واحدة عند التشغيل
MODEL_NAME = "aubmindlab/aragpt2-base"  # نموذج مخصص للعربية
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# 3. تعريف دالة التوليد
def generate_text(prompt: str, max_length=200):
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=max_length,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# 4. واجهة API باستخدام FastAPI
class Request(BaseModel):
    text: str
    max_length: int = 200

@app.post("/generate")
async def api_generate(request: Request):
    result = generate_text(request.text, request.max_length)
    return {"generated_text": result}

# 5. واجهة المستخدم باستخدام Gradio
def gradio_interface(prompt):
    return generate_text(prompt)

ui = gr.Interface(
    fn=gradio_interface,
    inputs=gr.Textbox(lines=3, placeholder="اكتب سؤالك هنا...", label="المدخلات"),
    outputs=gr.Textbox(label="الإجابة"),
    title="نموذج ذكاء اصطناعي للغة العربية",
    description="نموذج توليد نصوص بالعربية مع واجهة API"
)

# 6. دمج الواجهات
app = gr.mount_gradio_app(app, ui, path="/ui")

# 7. لتشغيل التطبيق محلياً (إختياري)
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)