Update app.py

app.py

@@ -1,23 +1,35 @@
 import os
 os.system("pip install -q accelerate")
 
-import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
-from threading import Thread
+import gradio as gr
 import spaces
+from threading import Thread
+
+# ======================================================
+# إعداد الموديل الخاص بك
+# ======================================================
+MODEL_ID = "anaspro/iraqi-kashif-2b"  # ✅ موديلك الخاص
+HF_TOKEN = os.environ.get("HF_TOKEN")  # من Secrets في Space
 
-MODEL_NAME = "anaspro/iraqi-kashif-2b"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+if HF_TOKEN:
+    print("🔐 HF_TOKEN detected, using authenticated download...")
+else:
+    print("⚠️ Warning: HF_TOKEN not found. Public model access only (may fail if private).")
 
+# ======================================================
+# تحميل النموذج داخل ZeroGPU Worker (لتفادي CUDA init)
+# ======================================================
 @spaces.GPU
 def load_model():
-    print("🔄 Loading model and tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    print("🔄 Loading Iraqi Kashif 2B model...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
     model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        dtype=torch.float16,
+        MODEL_ID,
+        dtype=torch.bfloat16,
         device_map="auto",
+        token=HF_TOKEN,
     )
     model.eval()
     print("✅ Model loaded successfully!")
@@ -25,57 +37,86 @@ def load_model():
 
 tokenizer, model = load_model()
 
-@spaces.GPU
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-    messages.extend(history)
+# ======================================================
+# دالة الدردشة (Streaming)
+# ======================================================
+@spaces.GPU(duration=60)
+def chat(message, history):
+    """
+    نموذج Iraqi Kashif 2B مدرب للهجة العراقية والفصحى
+    """
+    # بناء سياق المحادثة
+    messages = []
+    for user_msg, bot_msg in history:
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    prompt = tokenizer.apply_chat_template(
+    # تطبيق القالب من chat_template.jinja
+    input_ids = tokenizer.apply_chat_template(
         messages,
-        tokenize=False,
-        add_generation_prompt=True,
-    )
+        return_tensors="pt",
+        add_generation_prompt=True
+    ).to(model.device)
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    # Streamer للبث الحي أثناء التوليد
+    streamer = TextIteratorStreamer(
+        tokenizer,
+        skip_prompt=True,
+        skip_special_tokens=True
+    )
 
     generation_kwargs = dict(
-        **inputs,
+        input_ids=input_ids,
         streamer=streamer,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
         do_sample=True,
+        repetition_penalty=1.1,
     )
 
+    # تشغيل التوليد في Thread منفصل
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    response = ""
+    partial_text = ""
     for new_text in streamer:
-        response += new_text
-        yield response
+        partial_text += new_text
+        yield partial_text
+
+    thread.join()
 
+# ======================================================
+# واجهة Gradio
+# ======================================================
+demo = gr.ChatInterface(
+    fn=chat,
+    title="🇮🇶 Iraqi Kashif 2B - Chat Demo",
+    description="""
+    **نموذج Iraqi Kashif 2B**
 
-# ✅ لا نستخدم gr.get_state() إطلاقًا، فقط Gradio ChatInterface العادي
-chatbot = gr.ChatInterface(
-    fn=respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="أنت مساعد ذكي باللهجة العراقية.", label="System message"),
-        gr.Slider(minimum=32, maximum=1024, value=256, step=8, label="Max tokens"),
-        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-p"),
+    🧠 **مميزات النموذج:**
+    - مدرب على اللهجة العراقية + العربية الفصحى  
+    - يدعم المحادثة الطبيعية والفهم السياقي  
+    - مناسب لمشاريع الـ Chatbots و Call Centers  
+    
+    ⚙️ **ملاحظات:**
+    - يعمل على GPU باستخدام ZeroGPU
+    - يستخدم `bfloat16` لتقليل استهلاك VRAM
+    """,
+    examples=[
+        ["شلونك اليوم؟"],
+        ["اشتعني كلمة دلع؟"],
+        ["احجيلي نكتة عراقية"],
+        ["ما الفرق بين الإنترنت والأنتربيت؟"],
+        ["الشركة تفتح من 9 للـ 5، شنو أوقات الدوام؟"],
     ],
+    theme=gr.themes.Soft(),
 )
 
+# ======================================================
+# إطلاق التطبيق
+# ======================================================
 if __name__ == "__main__":
-    chatbot.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)