import os
import threading
import time
import torch
import inspect
import gc
from flask import Flask, request, jsonify, Response
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, BitsAndBytesConfig, TrainerCallback
from peft import LoraConfig, get_peft_model, TaskType
from datasets import Dataset
from trl import SFTTrainer, SFTConfig
app = Flask(__name__)
# ── Global state ─────────────────────────────────────────────────
_lock = threading.Lock()
state = {
"phase": "idle",
"progress": 0,
"message": "",
"eta": 0,
"model": None,
"tokenizer": None,
"model_name": None,
"error": None,
"download_start": None,
"train_start": None,
"total_steps": 0,
"current_step": 0,
"loss": None,
}
MODELS = [
{"id": "google/gemma-4-E2B-it", "name": "Gemma 4 E2B", "desc": "2B parameters • Google • Instruction-tuned", "size": "~4.5 GB", "ram": "8 GB+"},
{"id": "google/gemma-2-2b-it", "name": "Gemma 2 2B", "desc": "2B parameters • Google • Instruction-tuned", "size": "~5 GB", "ram": "8 GB+"},
{"id": "Qwen/Qwen2.5-1.5B-Instruct", "name": "Qwen 2.5 1.5B", "desc": "1.5B parameters • Alibaba • Instruction-tuned", "size": "~3 GB", "ram": "6 GB+"},
]
# ── الواجهة HTML مدمجة مباشرة ─────────────────────────────────────
HTML = """
ZYZ Fine-Tuner
⚡ ZYZ Fine-Tuner
The ZYZ Studio — ضبط دقيق لنماذج LLM بـ LoRA
الحالة idle
—
الوقت: 0ث
المتبقي: —
الخطوة: —
Loss: —
اختر النموذج
بيانات التدريب (JSON)
أدخل قائمة JSON من المحادثات بصيغة messages[]
تجربة النموذج
"""
# ── Routes ────────────────────────────────────────────────────────
@app.route("/")
def index():
return Response(HTML, mimetype="text/html")
@app.route("/api/models")
def get_models():
return jsonify(MODELS)
@app.route("/api/state")
def get_state():
now = time.time()
elapsed = 0
with _lock:
if state["download_start"]:
elapsed = int(now - state["download_start"])
if state["phase"] == "training" and state["train_start"]:
elapsed = int(now - state["train_start"])
return jsonify({
"phase": state["phase"],
"progress": state["progress"],
"message": state["message"],
"elapsed": elapsed,
"eta": state["eta"],
"loss": state["loss"],
"total_steps": state["total_steps"],
"current_step": state["current_step"],
"error": state["error"],
"model_name": state["model_name"],
})
@app.route("/api/download", methods=["POST"])
def start_download():
with _lock:
if state["phase"] not in ("idle", "done", "error"):
return jsonify({"ok": False, "msg": "العملية قيد التشغيل"}), 400
model_id = request.json.get("model_id", "google/gemma-4-E2B-it")
_reset_state()
state["phase"] = "downloading"
state["model_name"] = model_id
state["download_start"] = time.time()
threading.Thread(target=_download_model, args=(model_id,), daemon=True).start()
return jsonify({"ok": True})
@app.route("/api/train", methods=["POST"])
def start_train():
with _lock:
if state["phase"] != "ready":
return jsonify({"ok": False, "msg": "النموذج غير محمّل بعد"}), 400
dataset_json = request.json.get("dataset", [])
if not dataset_json:
return jsonify({"ok": False, "msg": "البيانات فارغة"}), 400
state["phase"] = "training"
state["train_start"] = time.time()
state["progress"] = 0
state["message"] = "جاري بدء التدريب..."
threading.Thread(target=_train_model, args=(dataset_json,), daemon=True).start()
return jsonify({"ok": True})
@app.route("/api/infer", methods=["POST"])
def infer():
with _lock:
model_ready = state["model"] is not None
if not model_ready:
return jsonify({"ok": False, "msg": "النموذج غير محمّل"}), 400
prompt = request.json.get("prompt", "")
if not prompt:
return jsonify({"ok": False, "msg": "الرسالة فارغة"}), 400
try:
tok = state["tokenizer"]
mdl = state["model"]
inputs = tok(prompt, return_tensors="pt").to(mdl.device)
with torch.no_grad():
out = mdl.generate(
**inputs,
max_new_tokens=256,
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tok.eos_token_id,
)
text = tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
return jsonify({"ok": True, "response": text})
except Exception as e:
return jsonify({"ok": False, "msg": str(e)}), 500
@app.route("/api/save", methods=["POST"])
def save_model():
with _lock:
model_ready = state["model"] is not None
if not model_ready:
return jsonify({"ok": False, "msg": "لا يوجد نموذج"}), 400
try:
save_dir = "/data/finetuned_model"
os.makedirs(save_dir, exist_ok=True)
state["model"].save_pretrained(save_dir)
state["tokenizer"].save_pretrained(save_dir)
return jsonify({"ok": True, "path": save_dir})
except Exception as e:
# fallback إلى /tmp إذا فشل /data
try:
state["model"].save_pretrained("/tmp/finetuned_model")
state["tokenizer"].save_pretrained("/tmp/finetuned_model")
return jsonify({"ok": True, "path": "/tmp/finetuned_model"})
except Exception as e2:
return jsonify({"ok": False, "msg": str(e2)}), 500
@app.route("/api/reset", methods=["POST"])
def reset():
with _lock:
# تحرير الذاكرة أولاً قبل إعادة الضبط
old_model = state.get("model")
old_tokenizer = state.get("tokenizer")
del old_model
del old_tokenizer
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
with _lock:
_reset_state()
return jsonify({"ok": True})
# ── Helpers ───────────────────────────────────────────────────────
def _unwrap_gemma4_linears(model):
"""استبدال Gemma4ClippableLinear بـ nn.Linear قبل تطبيق LoRA."""
to_replace = [
name for name, mod in model.named_modules()
if type(mod).__name__ == "Gemma4ClippableLinear"
]
for full_name in to_replace:
parts = full_name.split(".")
parent = model
for p in parts[:-1]:
parent = getattr(parent, p)
wrapper = getattr(parent, parts[-1])
setattr(parent, parts[-1], wrapper.linear)
print(f"[ZYZ] unwrapped {len(to_replace)} Gemma4ClippableLinear → nn.Linear")
return model
def _get_lora_targets(model_id: str) -> list:
m = model_id.lower()
if "gemma" in m or "qwen" in m:
return ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
return ["q_proj", "v_proj"]
def _calc_total_steps(num_examples: int, epochs: int, batch_size: int, grad_accum: int) -> int:
"""حساب عدد الخطوات الفعلي بدقة."""
steps_per_epoch = max(num_examples // (batch_size * grad_accum), 1)
return steps_per_epoch * epochs
def _build_sft_trainer(mdl, tok, ds, callbacks):
"""بناء SFTTrainer متوافق مع جميع نسخ trl."""
sig_trainer = inspect.signature(SFTTrainer.__init__).parameters
sig_config = inspect.signature(SFTConfig.__init__).parameters
use_gpu = torch.cuda.is_available()
# fp16/bf16 يشتغلان فقط على GPU — على CPU يجب إيقافهما تماماً
use_fp16 = use_gpu and torch.cuda.is_bf16_supported() is False
use_bf16 = use_gpu and torch.cuda.is_bf16_supported()
common = dict(
output_dir="/tmp/zyz_train_output",
num_train_epochs=3,
per_device_train_batch_size=1,
gradient_accumulation_steps=4,
warmup_steps=2,
learning_rate=2e-4,
fp16=use_fp16,
bf16=use_bf16,
logging_steps=1,
save_strategy="no",
report_to="none",
optim="paged_adamw_8bit" if use_gpu else "adamw_torch",
)
config_kwargs = common.copy()
trainer_kwargs = {}
if "max_seq_length" in sig_config:
config_kwargs["max_seq_length"] = 512
elif "max_seq_length" in sig_trainer:
trainer_kwargs["max_seq_length"] = 512
if "dataset_text_field" in sig_config:
config_kwargs["dataset_text_field"] = "text"
elif "dataset_text_field" in sig_trainer:
trainer_kwargs["dataset_text_field"] = "text"
train_cfg = SFTConfig(**config_kwargs)
trainer_kwargs["model"] = mdl
trainer_kwargs["train_dataset"] = ds
trainer_kwargs["args"] = train_cfg
trainer_kwargs["callbacks"] = callbacks
if "processing_class" in sig_trainer:
trainer_kwargs["processing_class"] = tok
elif "tokenizer" in sig_trainer:
trainer_kwargs["tokenizer"] = tok
return SFTTrainer(**trainer_kwargs)
# ── Background workers ────────────────────────────────────────────
def _download_model(model_id: str):
try:
with _lock:
state["message"] = f"جاري تحميل {model_id}..."
state["progress"] = 5
hf_token = os.environ.get("HF_TOKEN", None)
bnb_config = None
if torch.cuda.is_available():
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)
with _lock:
state["progress"] = 15
state["message"] = "تحميل الـ tokenizer..."
tokenizer = AutoTokenizer.from_pretrained(
model_id, token=hf_token, trust_remote_code=True,
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
with _lock:
state["progress"] = 40
state["message"] = "تحميل أوزان النموذج... (قد يستغرق بضع دقائق)"
# على CPU: float32 إلزامي — bfloat16/float16 لا يدعمهما CPU بشكل كامل
if torch.cuda.is_available():
load_dtype = None # bnb_config يتحكم في dtype على GPU
else:
load_dtype = torch.float32
model = AutoModelForCausalLM.from_pretrained(
model_id,
token=hf_token,
quantization_config=bnb_config,
device_map="auto" if torch.cuda.is_available() else "cpu",
trust_remote_code=True,
torch_dtype=load_dtype,
)
with _lock:
state["progress"] = 85
state["message"] = "تهيئة LoRA..."
model = _unwrap_gemma4_linears(model)
lora_cfg = LoraConfig(
task_type=TaskType.CAUSAL_LM,
r=8,
lora_alpha=16,
lora_dropout=0.05,
target_modules=_get_lora_targets(model_id),
)
model = get_peft_model(model, lora_cfg)
model.print_trainable_parameters()
with _lock:
state["model"] = model
state["tokenizer"] = tokenizer
state["progress"] = 100
state["message"] = "✅ تم تحميل النموذج بنجاح!"
state["phase"] = "ready"
except Exception as e:
with _lock:
state["phase"] = "error"
state["error"] = str(e)
state["message"] = f"❌ خطأ في التحميل: {e}"
def _train_model(dataset_json: list):
try:
with _lock:
tok = state["tokenizer"]
mdl = state["model"]
texts = []
for item in dataset_json:
conv = ""
for m in item.get("messages", []):
if m["role"] == "user":
conv += f"user\n{m['content']}\n"
elif m["role"] == "assistant":
conv += f"model\n{m['content']}\n"
if conv.strip():
texts.append(conv.strip())
ds = Dataset.from_dict({"text": texts})
# حساب الخطوات الفعلية
total_steps = _calc_total_steps(
num_examples=len(texts),
epochs=3,
batch_size=1,
grad_accum=4,
)
total_steps = max(total_steps, 1)
with _lock:
state["total_steps"] = total_steps
class _ProgressCB(TrainerCallback):
def __init__(self):
self.step = 0
def on_step_end(self, args, trainer_state, control, **kwargs):
self.step += 1
with _lock:
state["progress"] = min(int(self.step / total_steps * 100), 99)
state["current_step"] = self.step
elapsed = time.time() - state["train_start"]
state["eta"] = int((total_steps - self.step) * (elapsed / self.step)) if self.step else 0
if trainer_state.log_history:
last = trainer_state.log_history[-1]
if "loss" in last:
state["loss"] = round(last["loss"], 4)
state["message"] = f"التدريب... الخطوة {self.step}/{total_steps}"
trainer = _build_sft_trainer(mdl, tok, ds, [_ProgressCB()])
trainer.train()
with _lock:
state["model"] = trainer.model
state["progress"] = 100
state["phase"] = "done"
state["message"] = "✅ اكتمل التدريب بنجاح!"
except Exception as e:
with _lock:
state["phase"] = "error"
state["error"] = str(e)
state["message"] = f"❌ خطأ في التدريب: {e}"
def _reset_state():
"""يجب استدعاؤها داخل _lock."""
for k, v in [
("phase", "idle"), ("progress", 0), ("message", ""), ("eta", 0),
("model", None), ("tokenizer", None), ("model_name", None),
("error", None), ("download_start", None), ("train_start", None),
("total_steps", 0), ("current_step", 0), ("loss", None),
]:
state[k] = v
if __name__ == "__main__":
port = int(os.environ.get("PORT", 7860))
app.run(host="0.0.0.0", port=port, debug=False)