app.py

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

# ============================================================
# Load Model and Tokenizer
# ============================================================

MODEL_ID = "azeddinShr/LFM2-1.2B-RAG-ARABIC-AdaLoRA"

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
print("✅ Model loaded successfully!")

# ============================================================
# Generation Function
# ============================================================

def generate_answer(context, question, max_tokens=150, temperature=0.0):
    """
    Generate answer from context and question
    
    Args:
        context: Background information (Arabic text)
        question: Question to answer (Arabic text)
        max_tokens: Maximum length of generated answer
        temperature: Sampling temperature (0.0 = deterministic)
    
    Returns:
        Generated answer
    """
    # Validate inputs
    if not context or not context.strip():
        return "⚠️ الرجاء إدخال السياق (Please provide context)"
    
    if not question or not question.strip():
        return "⚠️ الرجاء إدخال السؤال (Please provide a question)"
    
    try:
        # Format prompt
        prompt = f"استخدم السياق التالي للإجابة على السؤال:\n\n{context}\n\nالسؤال: {question}"
        
        # Tokenize
        messages = [{"role": "user", "content": prompt}]
        input_ids = tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(model.device)
        
        # Generate
        with torch.no_grad():
            outputs = model.generate(
                input_ids,
                max_new_tokens=int(max_tokens),
                temperature=float(temperature) if temperature > 0 else 0.0,
                do_sample=temperature > 0,
                pad_token_id=tokenizer.eos_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        
        # Decode
        answer = tokenizer.decode(
            outputs[0][input_ids.shape[1]:], 
            skip_special_tokens=True
        ).strip()
        
        return answer
    
    except Exception as e:
        return f"❌ Error: {str(e)}"

# ============================================================
# Example Questions
# ============================================================

examples = [
    [
        "جمال أحمد حمزة خاشقجي صحفي وإعلامي سعودي، ولد في 13 أكتوبر 1958 في المدينة المنورة. رأس عدّة مناصب لعدد من الصحف في السعودية، وتقلّد منصب مستشار، كما أنّه مدير عام قناة العرب الإخبارية سابقًا.",
        "من هو جمال خاشقجي؟"
    ],
    [
        "نيوم هو مشروع ضخم في شمال غرب المملكة العربية السعودية أطلقه ولي العهد الأمير محمد بن سلمان في أكتوبر 2017. تبلغ المساحة الإجمالية المخططة لنيوم 26,500 كيلومتر مربع بتكلفة تقدر بـ 500 مليار دولار.",
        "ما هي المساحة الإجمالية لمشروع نيوم؟"
    ],
    [
        "نيوم هو مشروع ضخم في شمال غرب المملكة العربية السعودية أطلقه ولي العهد الأمير محمد بن سلمان في أكتوبر 2017. تبلغ المساحة الإجمالية المخططة لنيوم 26,500 كيلومتر مربع بتكلفة تقدر بـ 500 مليار دولار. يتضمن المشروع مناطق متعددة منها ذا لاين (The Line)، وهي مدينة خطية مستقبلية، بالإضافة إلى سندالة وأوكساجون وتروجينا. من المخطط أن تكتمل أجزاء رئيسية من المشروع بحلول عام 2030 كجزء من رؤية السعودية 2030.",
        "ما هي المساحة الإجمالية المخططة لمشروع نيوم؟"
    ],
    [
        "في فبراير 2025، وقعت نيوم اتفاقية شراكة مع شركة داتافولت السعودية لتصميم مصنع ضخم للذكاء الاصطناعي سيكون بمثابة مركز بيانات. بموجب الاتفاقية، سيتم تنفيذ المشروع في أوكساجون، مدينة الصناعات النظيفة والمتقدمة التابعة لنيوم، على مراحل. ستشهد المرحلة الأولى استثماراً أولياً قدره 5 مليارات دولار، ومن المتوقع أن يكون جاهزاً للعمل بحلول عام 2028.",
        "كم تبلغ قيمة الاستثمار الأولي لمصنع الذكاء الاصطناعي في نيوم؟"
    ]
]

# ============================================================
# Gradio Interface
# ============================================================

with gr.Blocks(theme=gr.themes.Soft(), title="Arabic QA System") as demo:
    
    gr.Markdown(
        """
        # 🤖 نظام الأسئلة والأجوبة العربي
        # Arabic Question Answering System
        
        Fine-tuned LFM2-1.2B-RAG model for Arabic extractive question answering.
        Provide context in Arabic and ask a question to get an answer extracted from the context.
        
        **Model:** [LiquidAI/LFM2-1.2B-RAG](https://huggingface.co/LiquidAI/LFM2-1.2B-RAG) fine-tuned with AdaLoRA
        """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            context_input = gr.Textbox(
                label="السياق / Context",
                placeholder="أدخل النص الذي يحتوي على المعلومات...\nEnter the text containing the information...",
                lines=8,
                rtl=True
            )
            
            question_input = gr.Textbox(
                label="السؤال / Question",
                placeholder="اطرح سؤالك هنا...\nAsk your question here...",
                lines=2,
                rtl=True
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                max_tokens = gr.Slider(
                    minimum=50,
                    maximum=500,
                    value=150,
                    step=10,
                    label="Max Tokens",
                    info="Maximum length of generated answer"
                )
                
                temperature = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.0,
                    step=0.1,
                    label="Temperature",
                    info="0 = deterministic, higher = more creative"
                )
            
            submit_btn = gr.Button("🔍 Get Answer / احصل على الإجابة", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            answer_output = gr.Textbox(
                label="الإجابة / Answer",
                lines=10,
                rtl=True,
                show_copy_button=True
            )
    
    gr.Examples(
        examples=examples,
        inputs=[context_input, question_input],
        label="📝 Example Questions / أمثلة"
    )
    
    gr.Markdown(
        """
        ---
        ### ℹ️ How to use / كيفية الاستخدام
        
        1. **Paste Context**: Add Arabic text containing information
        2. **Ask Question**: Write your question about the context
        3. **Get Answer**: The model will extract the answer from the context
        
        **Note:** This model is optimized for Modern Standard Arabic and extractive QA tasks.
        
        ### 🔗 Links
        - [Model Card](https://huggingface.co/YOUR_USERNAME/lfm2-1.2b-arabic-qa-adalora)
        - [Dataset (ARCD)](https://huggingface.co/datasets/hsseinmz/arcd)
        - [Base Model](https://huggingface.co/LiquidAI/LFM2-1.2B-RAG)
        """
    )
    
    # Connect button to function
    submit_btn.click(
        fn=generate_answer,
        inputs=[context_input, question_input, max_tokens, temperature],
        outputs=answer_output
    )

# ============================================================
# Launch
# ============================================================

if __name__ == "__main__":
    demo.launch()