Zgemma / app.py
THEZYZSTUDIO's picture
Update app.py
aa33c27 verified
Raw
History Blame Contribute Delete
25.2 kB
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 = """<!DOCTYPE html>
<html lang="ar" dir="rtl">
<head>
<meta charset="UTF-8"/>
<meta name="viewport" content="width=device-width, initial-scale=1.0"/>
<title>ZYZ Fine-Tuner</title>
<style>
*{box-sizing:border-box;margin:0;padding:0}
body{font-family:'Segoe UI',Tahoma,sans-serif;background:#0d0d1a;color:#e0e0ff;min-height:100vh;padding:20px}
h1{text-align:center;font-size:1.8rem;color:#a78bfa;margin-bottom:6px}
.sub{text-align:center;color:#6b7280;font-size:.85rem;margin-bottom:24px}
.card{background:#1a1a2e;border:1px solid #2d2d4e;border-radius:12px;padding:20px;margin-bottom:16px}
.card h2{font-size:1rem;color:#a78bfa;margin-bottom:14px}
.models{display:grid;grid-template-columns:repeat(auto-fill,minmax(220px,1fr));gap:10px}
.model-card{border:2px solid #2d2d4e;border-radius:8px;padding:12px;cursor:pointer;transition:.2s}
.model-card:hover{border-color:#7c3aed}
.model-card.selected{border-color:#a78bfa;background:#2d1b69}
.model-card .name{font-weight:700;font-size:.95rem;color:#e9d5ff}
.model-card .desc{font-size:.75rem;color:#9ca3af;margin-top:4px}
.model-card .meta{font-size:.72rem;color:#6b7280;margin-top:6px}
.meta span{background:#1e1b4b;border-radius:4px;padding:2px 6px;margin-left:4px}
.btn{border:none;border-radius:8px;padding:10px 22px;cursor:pointer;font-size:.9rem;font-weight:600;transition:.2s}
.btn-primary{background:linear-gradient(135deg,#7c3aed,#4f46e5);color:#fff}
.btn-primary:hover{opacity:.85}
.btn-primary:disabled{opacity:.4;cursor:not-allowed}
.btn-danger{background:#7f1d1d;color:#fca5a5}
.btn-danger:hover{background:#991b1b}
.btn-green{background:#065f46;color:#6ee7b7}
.btn-green:hover{background:#047857}
.progress-wrap{background:#2d2d4e;border-radius:8px;height:18px;overflow:hidden;margin:10px 0}
.progress-bar{height:100%;background:linear-gradient(90deg,#7c3aed,#4f46e5);transition:width .4s ease;border-radius:8px}
.status-msg{font-size:.85rem;color:#c4b5fd;margin-top:6px;min-height:20px}
.stats{display:flex;gap:16px;flex-wrap:wrap;margin-top:8px}
.stat{background:#0d0d1a;border-radius:6px;padding:6px 12px;font-size:.8rem}
.stat span{color:#a78bfa;font-weight:700}
textarea{width:100%;background:#0d0d1a;border:1px solid #2d2d4e;border-radius:8px;color:#e0e0ff;padding:12px;font-family:monospace;font-size:.82rem;resize:vertical;min-height:140px;margin-bottom:10px}
textarea:focus{outline:none;border-color:#7c3aed}
.infer-row{display:flex;gap:8px;align-items:flex-start}
.infer-input{flex:1}
.response-box{background:#0d0d1a;border:1px solid #2d2d4e;border-radius:8px;padding:12px;font-size:.85rem;min-height:60px;white-space:pre-wrap;color:#86efac;margin-top:10px}
.error-box{background:#1c0a0a;border:1px solid #7f1d1d;border-radius:8px;padding:12px;color:#fca5a5;font-size:.82rem;margin-top:8px}
.badge{display:inline-block;border-radius:6px;padding:2px 8px;font-size:.72rem;font-weight:700}
.badge-idle{background:#1e293b;color:#94a3b8}
.badge-downloading{background:#1e3a5f;color:#60a5fa}
.badge-ready{background:#064e3b;color:#34d399}
.badge-training{background:#3b1f00;color:#fbbf24}
.badge-done{background:#064e3b;color:#34d399}
.badge-error{background:#450a0a;color:#f87171}
.row-btns{display:flex;gap:8px;flex-wrap:wrap;margin-top:10px}
.hint{font-size:.75rem;color:#6b7280;margin-top:6px}
@media(max-width:480px){h1{font-size:1.3rem}.models{grid-template-columns:1fr}}
</style>
</head>
<body>
<h1>⚡ ZYZ Fine-Tuner</h1>
<p class="sub">The ZYZ Studio — ضبط دقيق لنماذج LLM بـ LoRA</p>
<!-- الحالة -->
<div class="card">
<h2>الحالة <span id="badge" class="badge badge-idle">idle</span></h2>
<div class="progress-wrap"><div class="progress-bar" id="pbar" style="width:0%"></div></div>
<div class="status-msg" id="msg">—</div>
<div class="stats">
<div class="stat">الوقت: <span id="elapsed">0</span>ث</div>
<div class="stat">المتبقي: <span id="eta">—</span></div>
<div class="stat">الخطوة: <span id="steps">—</span></div>
<div class="stat">Loss: <span id="loss">—</span></div>
</div>
<div id="err-box" class="error-box" style="display:none"></div>
</div>
<!-- اختيار النموذج -->
<div class="card">
<h2>اختر النموذج</h2>
<div class="models" id="models-list"></div>
<div class="row-btns">
<button class="btn btn-primary" id="btn-download" onclick="startDownload()">⬇️ تحميل النموذج</button>
<button class="btn btn-danger" id="btn-reset" onclick="doReset()">🔄 إعادة ضبط</button>
</div>
</div>
<!-- بيانات التدريب -->
<div class="card">
<h2>بيانات التدريب (JSON)</h2>
<textarea id="dataset" placeholder='[{"messages":[{"role":"user","content":"مرحباً"},{"role":"assistant","content":"أهلاً بك!"}]}]'></textarea>
<p class="hint">أدخل قائمة JSON من المحادثات بصيغة messages[]</p>
<div class="row-btns">
<button class="btn btn-primary" id="btn-train" onclick="startTrain()" disabled>🚀 ابدأ التدريب</button>
</div>
</div>
<!-- الاستدلال -->
<div class="card">
<h2>تجربة النموذج</h2>
<div class="infer-row">
<textarea class="infer-input" id="prompt" rows="3" placeholder="اكتب رسالتك هنا..."></textarea>
</div>
<div class="row-btns">
<button class="btn btn-primary" id="btn-infer" onclick="doInfer()" disabled>💬 أرسل</button>
<button class="btn btn-green" id="btn-save" onclick="doSave()" disabled>💾 حفظ النموذج</button>
</div>
<div class="response-box" id="response" style="display:none"></div>
</div>
<script>
let selectedModel = "google/gemma-4-E2B-it";
let pollInterval = null;
// تحميل قائمة النماذج
async function loadModels() {
const res = await fetch("/api/models");
const models = await res.json();
const el = document.getElementById("models-list");
el.innerHTML = "";
models.forEach(m => {
const div = document.createElement("div");
div.className = "model-card" + (m.id === selectedModel ? " selected" : "");
div.onclick = () => selectModel(m.id);
div.id = "mc-" + m.id.replace(/\//g,"_");
div.innerHTML = `<div class="name">${m.name}</div>
<div class="desc">${m.desc}</div>
<div class="meta"><span>${m.size}</span><span>RAM ${m.ram}</span></div>`;
el.appendChild(div);
});
}
function selectModel(id) {
selectedModel = id;
document.querySelectorAll(".model-card").forEach(c => c.classList.remove("selected"));
const el = document.getElementById("mc-" + id.replace(/\//g,"_"));
if(el) el.classList.add("selected");
}
async function startDownload() {
const res = await fetch("/api/download", {
method:"POST", headers:{"Content-Type":"application/json"},
body: JSON.stringify({model_id: selectedModel})
});
const d = await res.json();
if(!d.ok) return alert(d.msg);
startPoll();
}
async function startTrain() {
const raw = document.getElementById("dataset").value.trim();
if(!raw) return alert("أدخل بيانات التدريب أولاً");
let dataset;
try { dataset = JSON.parse(raw); } catch(e) { return alert("JSON غير صالح: " + e.message); }
const res = await fetch("/api/train", {
method:"POST", headers:{"Content-Type":"application/json"},
body: JSON.stringify({dataset})
});
const d = await res.json();
if(!d.ok) return alert(d.msg);
startPoll();
}
async function doInfer() {
const prompt = document.getElementById("prompt").value.trim();
if(!prompt) return;
const box = document.getElementById("response");
box.style.display = "block";
box.textContent = "...جاري التفكير";
const res = await fetch("/api/infer", {
method:"POST", headers:{"Content-Type":"application/json"},
body: JSON.stringify({prompt})
});
const d = await res.json();
box.textContent = d.ok ? d.response : "❌ " + d.msg;
}
async function doSave() {
const res = await fetch("/api/save", {method:"POST"});
const d = await res.json();
alert(d.ok ? "✅ تم الحفظ في: " + d.path : "❌ " + d.msg);
}
async function doReset() {
if(!confirm("إعادة ضبط كاملة؟ سيتم حذف النموذج المحمّل.")) return;
await fetch("/api/reset", {method:"POST"});
document.getElementById("response").style.display = "none";
stopPoll();
await updateUI();
}
function startPoll() {
if(pollInterval) return;
pollInterval = setInterval(updateUI, 1200);
updateUI();
}
function stopPoll() {
if(pollInterval) { clearInterval(pollInterval); pollInterval = null; }
}
async function updateUI() {
let s;
try {
const res = await fetch("/api/state");
s = await res.json();
} catch(e) { return; }
// شريط التقدم
document.getElementById("pbar").style.width = s.progress + "%";
document.getElementById("msg").textContent = s.message || "—";
// Badge
const badge = document.getElementById("badge");
badge.className = "badge badge-" + s.phase;
badge.textContent = s.phase;
// إحصاءات
document.getElementById("elapsed").textContent = s.elapsed + "ث";
document.getElementById("eta").textContent = s.eta ? s.eta + "ث" : "—";
document.getElementById("steps").textContent = (s.total_steps > 0) ? s.current_step + "/" + s.total_steps : "—";
document.getElementById("loss").textContent = s.loss !== null ? s.loss : "—";
// خطأ
const errBox = document.getElementById("err-box");
if(s.error) { errBox.style.display="block"; errBox.textContent = "❌ " + s.error; }
else { errBox.style.display="none"; }
// أزرار
const ready = (s.phase === "ready" || s.phase === "done");
const idle = (s.phase === "idle" || s.phase === "done" || s.phase === "error");
document.getElementById("btn-download").disabled = !idle;
document.getElementById("btn-train").disabled = !ready || s.phase === "training";
document.getElementById("btn-infer").disabled = s.model_name === null;
document.getElementById("btn-save").disabled = s.model_name === null;
// وقف الـ polling عند الانتهاء
if(s.phase === "done" || s.phase === "error" || s.phase === "idle") stopPoll();
}
loadModels();
updateUI();
</script>
</body>
</html>"""
# ── 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"<start_of_turn>user\n{m['content']}<end_of_turn>\n"
elif m["role"] == "assistant":
conv += f"<start_of_turn>model\n{m['content']}<end_of_turn>\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)