Update app.py

app.py

@@ -1,17 +1,21 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoTokenizer, TextIteratorStreamer
 import torch
 import gradio as gr
 import spaces
 from threading import Thread
 
+# نحاول استخدام Unsloth إذا متوفر
+try:
+    from unsloth import FastLanguageModel
+    HAS_UNSLOTH = True
+except ImportError:
+    HAS_UNSLOTH = False
+
 # ======================================================
 # إعدادات الموديل
 # ======================================================
 MODEL_ID = "unsloth/gemma-3-4b-it-unsloth-bnb-4bit"
 
-# ======================================================
-# البرومبت الافتراضي (متوافق مع أسلوب التدريب)
-# ======================================================
 SYSTEM_PROMPT = (
     "أنت مساعد ذكي تفهم اللهجة العراقية والعربية الفصحى. "
     "جاوب على الأسئلة بإيجاز ووضوح، بنفس لغة المستخدم. "
@@ -19,17 +23,30 @@ SYSTEM_PROMPT = (
 )
 
 # ======================================================
-# تحميل الموديل (مرة واحدة فقط)
+# تحميل الموديل (مع دعم Unsloth أو Transformers العادي)
 # ======================================================
-print("🔄 Loading Gemma 3 4B model...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-)
+print("🔄 Loading model:", MODEL_ID)
+
+if HAS_UNSLOTH:
+    print("🚀 Using Unsloth FastLanguageModel backend")
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        MODEL_ID,
+        dtype=torch.bfloat16,
+        load_in_4bit=True,
+        device_map="auto",
+    )
+else:
+    print("⚙️ Using standard Transformers backend")
+    from transformers import AutoModelForCausalLM
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        device_map="auto",
+    )
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
 model.eval()
-print("✅ Model ready!")
+print("✅ Model ready!\n")
 
 
 # ======================================================
@@ -37,7 +54,7 @@ print("✅ Model ready!")
 # ======================================================
 @spaces.GPU(duration=60)
 def chat(message, history):
-    # تحويل تاريخ المحادثة لصيغة مناسبة
+    # تحويل تاريخ المحادثة لصيغة messages
     messages = []
     for msg in history:
         if msg["role"] == "user":
@@ -48,7 +65,7 @@ def chat(message, history):
     # نضيف السؤال الحالي مع system prompt
     messages.append({"role": "user", "content": f"{SYSTEM_PROMPT}\n\nالسؤال: {message}"})
 
-    # تحويل الرسائل إلى input ids باستخدام chat template الخاص بـ Gemma
+    # تجهيز الإدخال عبر الـ chat template
     input_ids = tokenizer.apply_chat_template(
         messages,
         return_tensors="pt",
@@ -68,7 +85,7 @@ def chat(message, history):
         repetition_penalty=1.15,
     )
 
-    # توليد النص في thread منفصل للبث الحي
+    # تشغيل التوليد في Thread منفصل
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
@@ -86,12 +103,12 @@ def chat(message, history):
 demo = gr.ChatInterface(
     fn=chat,
     type="messages",
-    title="🇮🇶 Gemma 3 4B – Iraqi/Arabic Chat Demo",
+    title="🇮🇶 Gemma 3 4B – Iraqi Arabic Chat",
     description="""
 **نموذج Gemma 3 4B Instruct (Unsloth Optimized)**  
-مدرّب على العربية والفصحى ومهيّأ للحوار باللهجة العراقية.
+مدرّب على العربية والفصحى ومهيّأ للحوار باللهجة العراقية 🇮🇶
 
-🧠 جرّب:
+🧠 أمثلة جاهزة:
 - "شنو معنى إنكار معقول؟"
 - "اشرحلي شنو يعني تصعيد إداري"
 - "وضحلي الفرق بين الدليل والعبرة"