import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification, GPTNeoXForCausalLM

# تحميل النماذج
tokenizer_ara = AutoTokenizer.from_pretrained("aubmindlab/bert-base-arabertv02")
model_ara = AutoModelForSequenceClassification.from_pretrained("aubmindlab/bert-base-arabertv02")
tokenizer_gpt = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
model_gpt = GPTNeoXForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b")

# دالة للتنبؤ
def generate_text(input_text):
    input_ids_ara = tokenizer_ara(input_text, return_tensors="pt").input_ids
    outputs_ara = model_ara(input_ids_ara)
    input_ids_gpt = tokenizer_gpt(outputs_ara.logits.argmax(-1).item(), return_tensors="pt").input_ids
    outputs_gpt = model_gpt.generate(input_ids_gpt, max_length=100)
    generated_text = tokenizer_gpt.decode(outputs_gpt[0])
    return generated_text

# واجهة المستخدم
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=5, placeholder="أدخل نصًا عربيًا هنا..."),
    outputs=gr.Textbox(lines=5),
    title="توليد النصوص العربية باستخدام AraBERT و GPT-NeoX",
    description="أدخل نصًا عربيًا وستقوم هذه الأداة بتحليله باستخدام AraBERT ثم توليد نص جديد باستخدام GPT-NeoX."
)

iface.launch()