edits

app.py

@@ -1,20 +1,21 @@
-# ===============================
-# 1️⃣ استيراد المكتبات
-# ===============================
+# flask_lora_server.py
+# نسخة مُحسّنة من سكربت Flask لتشغيل موديل أساسي + LoRA + تدريب + توليد
+# تتضمن تحسينات: فحص bitsandbytes، تحميل وحفظ LoRA آمن، توليد آمن، وحماية أساسية (قابلة للتفعيل لاحقًا)
+
 import os
-os.environ["HF_HOME"] = "/tmp"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp"
-os.environ["PEFT_HOME"] = "/tmp"
+import json
+import threading
+import torch
+from flask import Flask, request, jsonify
+from flask_cors import CORS
 
+# Transformers / PEFT
 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, json
 from datasets import Dataset
 
 # ===============================
-# 2️⃣ إعداد المتغيرات
+# 1️⃣ إعداد المسارات والمتغيرات
 # ===============================
 MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
 FLASK_PORT = 7860
@@ -27,33 +28,46 @@ os.makedirs(SAVE_DIR, exist_ok=True)
 DATASET_FILE = os.path.join(DATA_DIR, "dataset.json")
 
 # ===============================
-# 3️⃣ تحديد الجهاز
+# 2️⃣ تحديد الجهاز ووجود bitsandbytes
 # ===============================
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print("🔹 Using device:", device)
 
+use_8bit = False
+if device == "cuda":
+    try:
+        import bitsandbytes  # noqa
+        use_8bit = True
+        print("🔸 bitsandbytes available: enabling 8-bit loading")
+    except Exception as e:
+        print("⚠️ bitsandbytes not available — will not use 8-bit:", e)
+
 # ===============================
-# 4️⃣ تحميل الموديل الأساسي بأمان
+# 3️⃣ تحميل الـ Tokenizer
 # ===============================
-print("🔄 Loading base model...")
+print("🔄 Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
+# ===============================
+# 4️⃣ تحميل الموديل الأساسي (مع تعامل آمن للـ 8-bit)
+# ===============================
+print("🔄 Loading base model...")
 try:
-    model = AutoModelForCausalLM.from_pretrained(
+    base_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
+        device_map="auto" if device == "cuda" else {"": "cpu"},
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        load_in_8bit=use_8bit,
+        trust_remote_code=True,
     )
-    model.to(device)
+    base_model.to(device)
     print("✅ Base model loaded!")
 except Exception as e:
     print("⚠️ Failed to load base model:", e)
     raise e
 
 # ===============================
-# 5️⃣ إعداد LoRA
+# 5️⃣ إعداد LoRA على الموديل
 # ===============================
 lora_config = LoraConfig(
     task_type=TaskType.CAUSAL_LM,
@@ -62,23 +76,31 @@ lora_config = LoraConfig(
     lora_dropout=0.05,
     target_modules=["q_proj", "v_proj"]
 )
-model = get_peft_model(model, lora_config)
-model.to(device)
-print("✅ LoRA model ready on device:", device)
+
+# Wrap base model with PEFT LoRA
+try:
+    model = get_peft_model(base_model, lora_config)
+    model.to(device)
+    print("✅ LoRA model ready on device:", device)
+except Exception as e:
+    print("⚠️ Failed to apply LoRA wrapper:", e)
+    raise
 
 # ===============================
-# 6️⃣ Auto load LoRA weights
+# 6️⃣ محاولة تحميل أوزان LoRA محفوظة (آمنة)
 # ===============================
-if os.path.exists(SAVE_DIR):
+if os.path.isdir(SAVE_DIR) and any(os.scandir(SAVE_DIR)):
     try:
-        base_model = AutoModelForCausalLM.from_pretrained(
+        print("🔄 Attempting to load LoRA from", SAVE_DIR)
+        # reload base model (fresh) and then load peft weights
+        base_model_for_peft = 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
+            device_map="auto" if device == "cuda" else {"": "cpu"},
+            torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+            load_in_8bit=use_8bit,
+            trust_remote_code=True,
         )
-        model = PeftModel.from_pretrained(base_model, SAVE_DIR)
+        model = PeftModel.from_pretrained(base_model_for_peft, SAVE_DIR, device_map="auto" if device == "cuda" else {"": "cpu"})
         model.to(device)
         print("✅ LoRA weights auto-loaded from disk!")
     except Exception as e:
@@ -95,7 +117,7 @@ CORS(app)
 trainer = None
 current_dataset = None
 
-# Fake DB
+# Fake DB — غيّر أو اربطه بقاعدة حقيقية لاحقًا
 users_db = {
     "1": {"name": "Ahmed Elsayed", "role": "Student"},
     "2": {"name": "Mona Ali", "role": "Lawyer"},
@@ -103,53 +125,61 @@ users_db = {
 }
 
 # ===============================
-# 7a️⃣ توليد نصوص
+# 7a️⃣ دالة توليد النصوص — محسّنة
 # ===============================
-@app.route("/generate", methods=["POST"])
+@app.route("/generate", methods=["POST"]) 
 def generate():
     try:
-        data = request.json
-        prompt = data.get("prompt", "").strip()
+        data = request.json or {}
+        prompt = (data.get("prompt") or "").strip()
         user_id = str(data.get("userId", "")).strip()
 
+        if not prompt:
+            return jsonify({"error": "Empty prompt"}), 400
+
         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}! أنا زين، مساعدك الذكي. برجاء طرح سؤال قانوني."}
-        #         ]
-        #     })
-
-        # Search in dataset
-        for item in current_dataset:
-            lines = item["text"].split("\n")
-            for line in lines:
-                if line.lower().startswith("user:") and prompt in line:
-                    response_messages = []
-                    for l in lines:
-                        if ":" in l:
-                            role, content = l.split(":", 1)
-                            response_messages.append({
-                                "role": role.strip().lower(),
-                                "content": content.strip()
-                            })
-                    response_messages.insert(0, {"role": "assistant", "content": f"أهلاً {user_name} 👋"})
-                    return jsonify({"messages": response_messages})
-
-        # Fallback: generate
-        inputs = tokenizer(prompt, return_tensors="pt").to(device)
+        # البحث في الداتا إذا كانت موجودة
+        if current_dataset is not None:
+            prompt_lower = prompt.lower()
+            for item in current_dataset:
+                text = item.get("text", "")
+                lines = text.split("\n")
+                for line in lines:
+                    line_lower = line.lower()
+                    if line_lower.startswith("user:") and prompt_lower in line_lower:
+                        response_messages = []
+                        for l in lines:
+                            if ":" in l:
+                                role, content = l.split(":", 1)
+                                response_messages.append({
+                                    "role": role.strip().lower(),
+                                    "content": content.strip()
+                                })
+                        response_messages.insert(0, {"role": "assistant", "content": f"أهلاً {user_name} 👋"})
+                        return jsonify({"messages": response_messages})
+
+        # Fallback: generate using النموذج
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True).to(device)
+        input_ids = inputs["input_ids"]
+        attention_mask = inputs.get("attention_mask", None)
+
         outputs = model.generate(
-            **inputs,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
             max_new_tokens=150,
             temperature=0.7,
-            do_sample=True
+            do_sample=True,
+            pad_token_id=(getattr(tokenizer, 'eos_token_id', None) or getattr(tokenizer, 'pad_token_id', None))
         )
-        reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        # اقتطاع جزء الـ prompt من الناتج للحصول على الجزء المولَّد فقط
+        generated_ids = outputs[0][input_ids.shape[-1]:]
+        reply = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
 
         return jsonify({
             "messages": [
@@ -162,15 +192,16 @@ def generate():
         return jsonify({"messages": [{"role": "assistant", "content": f"خطأ: {str(e)}"}]})
 
 # ===============================
-# 7b️⃣ تحميل الداتا
+# 7b️⃣ تحميل الداتا من ملف dataset.json
 # ===============================
+
 def load_dataset():
     global current_dataset
     if os.path.exists(DATASET_FILE):
         with open(DATASET_FILE, "r", encoding="utf-8") as f:
             try:
                 data = json.load(f)
-            except:
+            except Exception:
                 data = []
 
         if data:
@@ -183,8 +214,12 @@ def load_dataset():
                     formatted.append({"text": conversation.strip()})
                 elif "text" in item:
                     formatted.append({"text": item["text"]})
-            current_dataset = Dataset.from_list(formatted)
-            print(f"✅ Loaded {len(formatted)} records from dataset.json")
+            try:
+                current_dataset = Dataset.from_list(formatted)
+                print(f"✅ Loaded {len(formatted)} records from dataset.json")
+            except Exception as e:
+                print("⚠️ Failed to create Dataset:", e)
+                current_dataset = None
         else:
             current_dataset = None
     else:
@@ -193,7 +228,7 @@ def load_dataset():
 load_dataset()
 
 # ===============================
-# 7c️⃣ تدريب LoRA (Async)
+# 7c️⃣ تدريب LoRA (تشغيل في Thread منفصل)
 # ===============================
 @app.route("/train", methods=["POST"])
 def train():
@@ -203,7 +238,7 @@ def train():
         epochs = 1
         batch_size = 1
 
-        # 1️⃣ JSON مباشر
+        # JSON body مباشر أو dict يحتوي على data
         if request.is_json:
             body = request.get_json()
             if isinstance(body, list):
@@ -216,7 +251,6 @@ def train():
             epochs = int(body.get("epochs", 1)) if isinstance(body, dict) else 1
             batch_size = int(body.get("batch_size", 1)) if isinstance(body, dict) else 1
 
-        # 2️⃣ form-data ملف
         else:
             file = request.files.get("file")
             if file:
@@ -232,7 +266,7 @@ def train():
             with open(DATASET_FILE, "r", encoding="utf-8") as f:
                 try:
                     old_data = json.load(f)
-                except:
+                except Exception:
                     old_data = []
 
         combined = old_data + data
@@ -265,8 +299,9 @@ def train():
 
         current_dataset = Dataset.from_list(formatted)
 
+        # Tokenize — استخدم padding=True لتقليل استهلاك الذاكرة
         def tokenize_function(examples):
-            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256)
+            return tokenizer(examples["text"], truncation=True, padding=True, max_length=256)
 
         tokenized_dataset = current_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
 
@@ -276,18 +311,29 @@ def train():
             num_train_epochs=epochs,
             logging_steps=2,
             save_strategy="epoch",
-            fp16=torch.cuda.is_available(),
-            report_to=None
+            fp16=(device == "cuda"),
+            report_to=None,
         )
 
         trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_dataset)
 
         # تشغيل التدريب في Thread منفصل
         def run_training():
-            print("🔄 Training started...")
-            trainer.train()
-            trainer.save_model(SAVE_DIR)
-            print("✅ Training completed & model saved!")
+            try:
+                print("🔄 Training started...")
+                trainer.train()
+                # Trainer.save_model attempts to save the model; مع PEFT من الأفضل save_pretrained
+                try:
+                    trainer.save_model(SAVE_DIR)
+                except Exception as e:
+                    print("⚠️ trainer.save_model failed:", e)
+                try:
+                    model.save_pretrained(SAVE_DIR)
+                except Exception as e:
+                    print("⚠️ model.save_pretrained failed:", e)
+                print("✅ Training completed & model saved!")
+            except Exception as e:
+                print("❌ Training failed:", e)
 
         threading.Thread(target=run_training, daemon=True).start()
 
@@ -301,11 +347,10 @@ def train():
     except Exception as e:
         return jsonify({"error": str(e)})
 
-
 # ===============================
-# 8️⃣ تشغيل Flask
+# 8️⃣ Endpoint بسيط للفحص
 # ===============================
-@app.route("/", methods=["GET"])
+@app.route("/", methods=["GET"]) 
 def home():
     return jsonify({
         "status": "ok",
@@ -313,5 +358,9 @@ def home():
         "endpoints": ["/generate", "/train"]
     })
 
+# ===============================
+# 9️⃣ تشغيل Flask
+# ===============================
 if __name__ == "__main__":
+    # ملاحظة أمنيّة: لا تشغّل هذا السكربت مكشوفًا للعالم دون طبقة مصادقة أو proxy
     app.run(host="0.0.0.0", port=FLASK_PORT, threaded=True)