Spaces:
Sleeping
Sleeping
| # =============================== | |
| # 1️⃣ استيراد المكتبات | |
| # =============================== | |
| 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, json | |
| from datasets import Dataset | |
| # =============================== | |
| # 2️⃣ إعداد المتغيرات | |
| # =============================== | |
| MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" | |
| FLASK_PORT = 7860 | |
| SAVE_DIR = "/tmp/lora_model" | |
| DATA_DIR = "/tmp/data" | |
| os.makedirs(DATA_DIR, exist_ok=True) | |
| 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) | |
| # =============================== | |
| # 4️⃣ تحميل الموديل الأساسي بأمان | |
| # =============================== | |
| print("🔄 Loading base model...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| 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 | |
| # =============================== | |
| # 5️⃣ إعداد LoRA | |
| # =============================== | |
| lora_config = LoraConfig( | |
| task_type=TaskType.CAUSAL_LM, | |
| r=2, | |
| lora_alpha=8, | |
| 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) | |
| # =============================== | |
| # 6️⃣ Auto load LoRA weights | |
| # =============================== | |
| if os.path.exists(SAVE_DIR): | |
| try: | |
| 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 | |
| ) | |
| model = PeftModel.from_pretrained(base_model, SAVE_DIR) | |
| model.to(device) | |
| print("✅ LoRA weights auto-loaded from disk!") | |
| except Exception as e: | |
| print("⚠️ Failed to auto-load LoRA weights:", e) | |
| else: | |
| print("ℹ️ No saved LoRA weights found, running with base model.") | |
| # =============================== | |
| # 7️⃣ إعداد Flask API | |
| # =============================== | |
| app = Flask(__name__) | |
| CORS(app) | |
| trainer = None | |
| current_dataset = None | |
| # Fake DB | |
| users_db = { | |
| "1": {"name": "Ahmed Elsayed", "role": "Student"}, | |
| "2": {"name": "Mona Ali", "role": "Lawyer"}, | |
| "3": {"name": "Omar Hassan", "role": "Client"} | |
| } | |
| # =============================== | |
| # 7a️⃣ توليد نصوص | |
| # =============================== | |
| def generate(): | |
| try: | |
| data = request.json | |
| 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}! أنا زين، مساعدك الذكي. برجاء طرح سؤال قانوني."} | |
| ] | |
| }) | |
| # 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) | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=150, | |
| temperature=0.7, | |
| do_sample=True | |
| ) | |
| reply = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return jsonify({ | |
| "messages": [ | |
| {"role": "assistant", "content": f"أهلاً {user_name} 👋"}, | |
| {"role": "assistant", "content": reply} | |
| ] | |
| }) | |
| except Exception as e: | |
| return jsonify({"messages": [{"role": "assistant", "content": f"خطأ: {str(e)}"}]}) | |
| # =============================== | |
| # 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: | |
| data = [] | |
| if data: | |
| formatted = [] | |
| for item in data: | |
| if "messages" in item: | |
| conversation = "" | |
| for msg in item["messages"]: | |
| conversation += f"{msg['role'].capitalize()}: {msg['content']}\n" | |
| 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") | |
| else: | |
| current_dataset = None | |
| else: | |
| current_dataset = None | |
| load_dataset() | |
| # =============================== | |
| # 7c️⃣ تدريب LoRA (Async) | |
| # =============================== | |
| def train(): | |
| global trainer, current_dataset, model | |
| try: | |
| data = [] | |
| epochs = 1 | |
| batch_size = 1 | |
| # 1️⃣ JSON مباشر | |
| if request.is_json: | |
| body = request.get_json() | |
| if isinstance(body, list): | |
| data = body | |
| elif isinstance(body, dict) and "data" in body: | |
| 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 | |
| # 2️⃣ form-data ملف | |
| else: | |
| file = request.files.get("file") | |
| if file: | |
| data = json.load(file) | |
| epochs = int(request.form.get("epochs", 1)) | |
| batch_size = int(request.form.get("batch_size", 1)) | |
| 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: | |
| 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) | |
| # الحفظ | |
| with open(DATASET_FILE, "w", encoding="utf-8") as f: | |
| json.dump(unique_data, f, ensure_ascii=False, indent=2) | |
| # تجهيز Dataset للتدريب | |
| formatted = [] | |
| for item in unique_data: | |
| if "messages" in item: | |
| conversation = "" | |
| for msg in item["messages"]: | |
| conversation += f"{msg['role'].capitalize()}: {msg['content']}\n" | |
| formatted.append({"text": conversation.strip()}) | |
| elif "text" in item: | |
| formatted.append({"text": item["text"]}) | |
| current_dataset = Dataset.from_list(formatted) | |
| def tokenize_function(examples): | |
| return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=256) | |
| tokenized_dataset = current_dataset.map(tokenize_function, batched=True, remove_columns=["text"]) | |
| training_args = TrainingArguments( | |
| output_dir=SAVE_DIR, | |
| per_device_train_batch_size=batch_size, | |
| num_train_epochs=epochs, | |
| logging_steps=2, | |
| save_strategy="epoch", | |
| fp16=torch.cuda.is_available(), | |
| 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!") | |
| threading.Thread(target=run_training, daemon=True).start() | |
| return jsonify({ | |
| "status": "training started", | |
| "epochs": epochs, | |
| "batch_size": batch_size, | |
| "records": len(unique_data) | |
| }) | |
| except Exception as e: | |
| return jsonify({"error": str(e)}) | |
| # =============================== | |
| # 8️⃣ تشغيل Flask | |
| # =============================== | |
| def home(): | |
| return jsonify({ | |
| "status": "ok", | |
| "message": "Graduation Project API is running 🚀", | |
| "endpoints": ["/generate", "/train"] | |
| }) | |
| if __name__ == "__main__": | |
| app.run(host="0.0.0.0", port=FLASK_PORT, threaded=True) | |