edits

app.py

@@ -1,23 +1,22 @@
-# flask_lora_server.py
-# نسخة مُحسّنة من سكربت Flask لتشغيل موديل أساسي + LoRA + تدريب + توليد
-# تتضمن تحسينات: فحص bitsandbytes، تحميل وحفظ LoRA آمن، توليد آمن، وحماية أساسية (قابلة للتفعيل لاحقًا)
-
+# ===============================
+# 1️⃣ استيراد المكتبات
+# ===============================
 import os
-import json
-import threading
-import torch
+os.environ["HF_HOME"] = "/tmp"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp"
+os.environ["PEFT_HOME"] = "/tmp"
+
 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
+import threading, torch, json
 from datasets import Dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, BitsAndBytesConfig
+from peft import LoraConfig, get_peft_model, TaskType, PeftModel
 
 # ===============================
-# 1️⃣ إعداد المسارات والمتغيرات
+# 2️⃣ إعداد المتغيرات
 # ===============================
-MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
+MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"   # ✅ موديل 7B
 FLASK_PORT = 7860
 
 SAVE_DIR = "/tmp/lora_model"
@@ -28,80 +27,62 @@ os.makedirs(SAVE_DIR, exist_ok=True)
 DATASET_FILE = os.path.join(DATA_DIR, "dataset.json")
 
 # ===============================
-# 2️⃣ تحديد الجهاز ووجود bitsandbytes
+# 3️⃣ تحديد الجهاز
 # ===============================
 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)
-
 # ===============================
-# 3️⃣ تحميل الـ Tokenizer
+# 4️⃣ تحميل الموديل مع 4bit Quantization
 # ===============================
-print("🔄 Loading tokenizer...")
+print("🔄 Loading base model with 4-bit quantization...")
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,                 # ✅ تشغيل 4bit quantization
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# ===============================
-# 4️⃣ تحميل الموديل الأساسي (مع تعامل آمن للـ 8-bit)
-# ===============================
-print("🔄 Loading base model...")
 try:
-    base_model = AutoModelForCausalLM.from_pretrained(
+    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=use_8bit,
-        trust_remote_code=True,
+        device_map="auto",
+        quantization_config=bnb_config,
+        trust_remote_code=True
     )
-    base_model.to(device)
-    print("✅ Base model loaded!")
+    print("✅ Base model loaded in 4-bit!")
 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,
-    r=2,
-    lora_alpha=8,
-    lora_dropout=0.05,
+    r=8,
+    lora_alpha=16,
+    lora_dropout=0.1,
     target_modules=["q_proj", "v_proj"]
 )
-
-# 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
+model = get_peft_model(model, lora_config)
+print("✅ LoRA model ready on device:", device)
 
 # ===============================
-# 6️⃣ محاولة تحميل أوزان LoRA محفوظة (آمنة)
+# 6️⃣ Auto load LoRA weights
 # ===============================
-if os.path.isdir(SAVE_DIR) and any(os.scandir(SAVE_DIR)):
+if os.path.exists(SAVE_DIR):
     try:
-        print("🔄 Attempting to load LoRA from", SAVE_DIR)
-        # reload base model (fresh) and then load peft weights
-        base_model_for_peft = 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=use_8bit,
-            trust_remote_code=True,
+            device_map="auto",
+            quantization_config=bnb_config,
+            trust_remote_code=True
         )
-        model = PeftModel.from_pretrained(base_model_for_peft, SAVE_DIR, device_map="auto" if device == "cuda" else {"": "cpu"})
-        model.to(device)
+        model = PeftModel.from_pretrained(base_model, SAVE_DIR)
         print("✅ LoRA weights auto-loaded from disk!")
     except Exception as e:
         print("⚠️ Failed to auto-load LoRA weights:", e)
@@ -117,6 +98,7 @@ CORS(app)
 trainer = None
 current_dataset = None
 
+# Fake DB
 users_db = {
     "1": {"name": "Ahmed Elsayed", "role": "Student"},
     "2": {"name": "Mona Ali", "role": "Lawyer"},
@@ -124,61 +106,30 @@ users_db = {
 }
 
 # ===============================
-# 7a️⃣ دالة توليد النصوص — محسّنة
+# 7a️⃣ توليد نصوص
 # ===============================
-@app.route("/generate", methods=["POST"]) 
+@app.route("/generate", methods=["POST"])
 def generate():
     try:
-        data = request.json or {}
-        prompt = (data.get("prompt") or "").strip()
+        data = request.json
+        prompt = data.get("prompt", "").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 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)
-
+        # توليد نصوص من الموديل
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
         outputs = model.generate(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            max_new_tokens=150,
+            **inputs,
+            max_new_tokens=200,
             temperature=0.7,
-            do_sample=True,
-            pad_token_id=(getattr(tokenizer, 'eos_token_id', None) or getattr(tokenizer, 'pad_token_id', None))
+            do_sample=True
         )
-
-        # اقتطاع جزء الـ prompt من الناتج للحصول على الجزء المولَّد فقط
-        generated_ids = outputs[0][input_ids.shape[-1]:]
-        reply = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
+        reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
         return jsonify({
             "messages": [
@@ -191,18 +142,16 @@ def generate():
         return jsonify({"messages": [{"role": "assistant", "content": f"خطأ: {str(e)}"}]})
 
 # ===============================
-# 7b️⃣ تحميل الداتا من ملف dataset.json
+# 7b️⃣ تحميل الداتا
 # ===============================
-
 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 Exception:
+            except:
                 data = []
-
         if data:
             formatted = []
             for item in data:
@@ -213,12 +162,8 @@ def load_dataset():
                     formatted.append({"text": conversation.strip()})
                 elif "text" in item:
                     formatted.append({"text": item["text"]})
-            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
+            current_dataset = Dataset.from_list(formatted)
+            print(f"✅ Loaded {len(formatted)} records from dataset.json")
         else:
             current_dataset = None
     else:
@@ -227,7 +172,7 @@ def load_dataset():
 load_dataset()
 
 # ===============================
-# 7c️⃣ تدريب LoRA (تشغيل في Thread منفصل)
+# 7c️⃣ تدريب LoRA (Async)
 # ===============================
 @app.route("/train", methods=["POST"])
 def train():
@@ -237,7 +182,6 @@ def train():
         epochs = 1
         batch_size = 1
 
-        # JSON body مباشر أو dict يحتوي على data
         if request.is_json:
             body = request.get_json()
             if isinstance(body, list):
@@ -246,10 +190,8 @@ def train():
                 data = body["data"]
             else:
                 return jsonify({"error": "JSON format not recognized"}), 400
-
             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
-
         else:
             file = request.files.get("file")
             if file:
@@ -259,35 +201,12 @@ def train():
             else:
                 return jsonify({"error": "No file uploaded or JSON provided"}), 400
 
-        # دمج وإزالة التكرارات حسب محتوى الرسائل
-        old_data = []
-        if os.path.exists(DATASET_FILE):
-            with open(DATASET_FILE, "r", encoding="utf-8") as f:
-                try:
-                    old_data = json.load(f)
-                except Exception:
-                    old_data = []
-
-        combined = old_data + data
-        unique_data, seen_texts = [], set()
-        for item in combined:
-            if "messages" in item:
-                text_key = "\n".join([f"{m['role'].strip().lower()}:{m['content'].strip()}" for m in item["messages"]])
-            elif "text" in item:
-                text_key = item["text"].strip()
-            else:
-                continue
-            if text_key not in seen_texts:
-                seen_texts.add(text_key)
-                unique_data.append(item)
-
-        # الحفظ
+        # الحفظ في dataset.json
         with open(DATASET_FILE, "w", encoding="utf-8") as f:
-            json.dump(unique_data, f, ensure_ascii=False, indent=2)
+            json.dump(data, f, ensure_ascii=False, indent=2)
 
-        # تجهيز Dataset للتدريب
         formatted = []
-        for item in unique_data:
+        for item in data:
             if "messages" in item:
                 conversation = ""
                 for msg in item["messages"]:
@@ -298,9 +217,8 @@ def train():
 
         current_dataset = Dataset.from_list(formatted)
 
-        # Tokenize — استخدم padding=True لتقليل استهلاك الذاكرة
         def tokenize_function(examples):
-            return tokenizer(examples["text"], truncation=True, padding=True, max_length=256)
+            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)
 
         tokenized_dataset = current_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
 
@@ -308,31 +226,19 @@ def train():
             output_dir=SAVE_DIR,
             per_device_train_batch_size=batch_size,
             num_train_epochs=epochs,
-            logging_steps=2,
+            logging_steps=5,
             save_strategy="epoch",
-            fp16=(device == "cuda"),
-            report_to=None,
+            fp16=torch.cuda.is_available(),
+            report_to=None
         )
 
         trainer = Trainer(model=model, args=training_args, train_dataset=tokenized_dataset)
 
-        # تشغيل التدريب في Thread منفصل
         def run_training():
-            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)
+            print("🔄 Training started...")
+            trainer.train()
+            trainer.save_model(SAVE_DIR)
+            print("✅ Training completed & model saved!")
 
         threading.Thread(target=run_training, daemon=True).start()
 
@@ -340,26 +246,22 @@ def train():
             "status": "training started",
             "epochs": epochs,
             "batch_size": batch_size,
-            "records": len(unique_data)
+            "records": len(data)
         })
 
     except Exception as e:
         return jsonify({"error": str(e)})
 
 # ===============================
-# 8️⃣ Endpoint بسيط للفحص
+# 8️⃣ تشغيل Flask
 # ===============================
-@app.route("/", methods=["GET"]) 
+@app.route("/", methods=["GET"])
 def home():
     return jsonify({
         "status": "ok",
-        "message": "Graduation Project API is running 🚀",
+        "message": "Qwen-7B LoRA API is running 🚀",
         "endpoints": ["/generate", "/train"]
     })
 
-# ===============================
-# 9️⃣ تشغيل Flask
-# ===============================
 if __name__ == "__main__":
-    # ملاحظة أمنيّة: لا تشغّل هذا السكربت مكشوفًا للعالم دون طبقة مصادقة أو proxy
     app.run(host="0.0.0.0", port=FLASK_PORT, threaded=True)