Merge branch 'Ai' into main

app.py

@@ -5,17 +5,17 @@ import os
 os.environ["HF_HOME"] = "/tmp"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 os.environ["PEFT_HOME"] = "/tmp"
+
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 from peft import LoraConfig, get_peft_model, TaskType, PeftModel
 from flask import Flask, request, jsonify
 from flask_cors import CORS
-import threading, torch, os, json
+import threading, torch, json
 from datasets import Dataset
 
 # ===============================
 # 2️⃣ إعداد المتغيرات
 # ===============================
-
 MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
 FLASK_PORT = 7860
 
@@ -26,21 +26,34 @@ os.makedirs(SAVE_DIR, exist_ok=True)
 
 DATASET_FILE = os.path.join(DATA_DIR, "dataset.json")
 
+# ===============================
+# 3️⃣ تحديد الجهاز
+# ===============================
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print("🔹 Using device:", device)
 
 # ===============================
-# 3️⃣ تحميل الموديل الأساسي
+# 4️⃣ تحميل الموديل الأساسي بأمان
 # ===============================
 print("🔄 Loading base model...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-)
-print("✅ Base model loaded!")
+
+try:
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto" if device=="cuda" else {"": "cpu"},
+        torch_dtype=torch.float16 if device=="cuda" else torch.float32,
+        load_in_8bit=(device=="cuda"),
+        trust_remote_code=True
+    )
+    model.to(device)
+    print("✅ Base model loaded!")
+except Exception as e:
+    print("⚠️ Failed to load base model:", e)
+    raise e
 
 # ===============================
-# 4️⃣ إعداد LoRA
+# 5️⃣ إعداد LoRA
 # ===============================
 lora_config = LoraConfig(
     task_type=TaskType.CAUSAL_LM,
@@ -50,27 +63,31 @@ lora_config = LoraConfig(
     target_modules=["q_proj", "v_proj"]
 )
 model = get_peft_model(model, lora_config)
-print("✅ LoRA model ready!")
+model.to(device)
+print("✅ LoRA model ready on device:", device)
 
 # ===============================
-# 🟢 Auto load LoRA weights on startup
+# 6️⃣ Auto load LoRA weights
 # ===============================
 if os.path.exists(SAVE_DIR):
     try:
         base_model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
-            device_map="auto",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            device_map="auto" if device=="cuda" else {"": "cpu"},
+            torch_dtype=torch.float16 if device=="cuda" else torch.float32,
+            load_in_8bit=(device=="cuda"),
+            trust_remote_code=True
         )
         model = PeftModel.from_pretrained(base_model, SAVE_DIR)
-        print("✅ LoRA weights auto-loaded from disk! [Fine-tuned model running]")
+        model.to(device)
+        print("✅ LoRA weights auto-loaded from disk!")
     except Exception as e:
-        print("⚠️ Failed to auto-load LoRA weights:", str(e))
+        print("⚠️ Failed to auto-load LoRA weights:", e)
 else:
     print("ℹ️ No saved LoRA weights found, running with base model.")
 
 # ===============================
-# 5️⃣ إعداد Flask API
+# 7️⃣ إعداد Flask API
 # ===============================
 app = Flask(__name__)
 CORS(app)
@@ -78,9 +95,7 @@ CORS(app)
 trainer = None
 current_dataset = None
 
-# ===============================
-# Fake Database (مؤقت للتجربة)
-# ===============================
+# Fake DB
 users_db = {
     "1": {"name": "Ahmed Elsayed", "role": "Student"},
     "2": {"name": "Mona Ali", "role": "Lawyer"},
@@ -88,7 +103,7 @@ users_db = {
 }
 
 # ===============================
-# 5a️⃣ توليد نصوص (generate)
+# 7a️⃣ توليد نصوص
 # ===============================
 @app.route("/generate", methods=["POST"])
 def generate():
@@ -97,25 +112,20 @@ def generate():
         prompt = data.get("prompt", "").strip()
         user_id = str(data.get("userId", "")).strip()
 
-        # 🔍 التحقق من المستخدم
         user_info = users_db.get(user_id)
         if not user_info:
             return jsonify({"error": "User not found"}), 404
 
         user_name = user_info["name"]
 
-        # 📌 لو مفيش داتا حالياً
         if current_dataset is None:
             return jsonify({
                 "messages": [
-                    {
-                        "role": "assistant",
-                        "content": f"⚖️ أهلاً {user_name}! أنا زين، مساعدك الذكي. برجاء طرح سؤال قانوني."
-                    }
+                    {"role": "assistant", "content": f"⚖️ أهلاً {user_name}! أنا زين، مساعدك الذكي. برجاء طرح سؤال قانوني."}
                 ]
             })
 
-        # 🔎 البحث في الداتا
+        # Search in dataset
         for item in current_dataset:
             lines = item["text"].split("\n")
             for line in lines:
@@ -128,15 +138,11 @@ def generate():
                                 "role": role.strip().lower(),
                                 "content": content.strip()
                             })
-                    # 📌 إضافة اسم المستخدم في البداية
-                    response_messages.insert(0, {
-                        "role": "assistant",
-                        "content": f"أهلاً {user_name} 👋"
-                    })
+                    response_messages.insert(0, {"role": "assistant", "content": f"أهلاً {user_name} 👋"})
                     return jsonify({"messages": response_messages})
 
-        # 🟡 fallback (لو مفيش في الداتا)
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        # Fallback: generate
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
         outputs = model.generate(
             **inputs,
             max_new_tokens=150,
@@ -153,13 +159,10 @@ def generate():
         })
 
     except Exception as e:
-        return jsonify({
-            "messages": [
-                {"role": "assistant", "content": f"خطأ: {str(e)}"}
-            ]
-        })
+        return jsonify({"messages": [{"role": "assistant", "content": f"خطأ: {str(e)}"}]})
+
 # ===============================
-# 🟢 تحميل الداتا من الملف عند البداية
+# 7b️⃣ تحميل الداتا
 # ===============================
 def load_dataset():
     global current_dataset
@@ -190,7 +193,7 @@ def load_dataset():
 load_dataset()
 
 # ===============================
-# 5b️⃣ تدريب LoRA (Async) + حفظ
+# 7c️⃣ تدريب LoRA (Async)
 # ===============================
 @app.route("/train", methods=["POST"])
 def train():
@@ -200,14 +203,11 @@ def train():
         epochs = 1
         batch_size = 1
 
-        # 📌 1) JSON request
         if request.is_json:
             body = request.get_json()
             data = body.get("data", [])
             epochs = int(body.get("epochs", 1))
             batch_size = int(body.get("batch_size", 1))
-
-        # 📌 2) Form-data request
         else:
             epochs = int(request.form.get("epochs", 1))
             batch_size = int(request.form.get("batch_size", 1))
@@ -217,7 +217,7 @@ def train():
             else:
                 return jsonify({"error": "No file uploaded"}), 400
 
-        # دمج وإزالة التكرارات
+        # Merge & remove duplicates
         old_data = []
         if os.path.exists(DATASET_FILE):
             with open(DATASET_FILE, "r", encoding="utf-8") as f:
@@ -234,11 +234,11 @@ def train():
                 seen.add(key)
                 unique_data.append(item)
 
-        # الحفظ
+        # Save dataset
         with open(DATASET_FILE, "w", encoding="utf-8") as f:
             json.dump(unique_data, f, ensure_ascii=False, indent=2)
 
-        # تجهيز الداتا
+        # Prepare dataset
         formatted = []
         for item in unique_data:
             if "messages" in item:
@@ -282,10 +282,8 @@ def train():
         return jsonify({"error": str(e)})
 
 # ===============================
-# 6️⃣ تشغيل Flask (بدون ngrok)
+# 8️⃣ تشغيل Flask
 # ===============================
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=FLASK_PORT, threaded=True)
 @app.route("/", methods=["GET"])
 def home():
     return jsonify({
@@ -293,3 +291,6 @@ def home():
         "message": "Graduation Project API is running 🚀",
         "endpoints": ["/generate", "/train"]
     })
+
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=FLASK_PORT, threaded=True)