Spaces:

THEZYZSTUDIO
/

Other2

Sleeping

Other2 / app.py
Update app.py
5e96ca2 verified 22 days ago
40 kB
import os
import re
import json
import threading
import traceback
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer

MODEL_ID = "GnLOLot/MiniCPM5-2B-Claude-Fable5-1-Thinking-Agentic"

# Qwen3 فيه ميزة "enable_thinking" مدمجة داخل قالب المحادثة نفسه، أما
# Qwen2.5-Coder (ومعظم عائلة Qwen2.5) فليس لديه هذا المعامل في القالب.
# نضبط هذا العلم هنا مرة واحدة فقط، ونستخدمه بدل استدعاء enable_thinking
# مباشرة على tokenizer.apply_chat_template في أي مكان بالملف، حتى يبقى
# الكود يعمل مع أي عائلة Qwen (3 أو 2.5) بدون تعديل لاحقاً.
MODEL_SUPPORTS_NATIVE_THINKING = "qwen3" in MODEL_ID.lower()

MAX_CONTEXT_TOKENS = 20000
MAX_OUTPUT_TOKENS = 4096


def _detect_cpus():
    """عدد الأنوية الفعلي المسموح للحاوية (مهم جداً على HF Spaces:
    torch يرى أنوية الجهاز كله فيُبطئ نفسه بكثرة الخيوط)."""
    try:
        with open("/sys/fs/cgroup/cpu.max") as f:
            q, per = f.read().split()
            if q != "max":
                return max(1, int(int(q) / int(per)))
    except Exception:
        pass
    try:
        return len(os.sched_getaffinity(0))
    except Exception:
        return os.cpu_count() or 2


CPU_THREADS = int(os.environ.get("CPU_THREADS", _detect_cpus()))
torch.set_num_threads(CPU_THREADS)
print(f"عدد خيوط CPU المستخدمة: {CPU_THREADS}")

print(f"تحميل النموذج {MODEL_ID} ...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side="left")
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    dtype=torch.float32,   # float32 أسرع وأكثر استقراراً من bf16 على CPU
    attn_implementation="sdpa",
)
model.eval()

LORA_DIR = os.environ.get("OUTPUT_DIR", "ddgoai-lora")


def _lora_matches_base():
    try:
        with open(os.path.join(LORA_DIR, "adapter_config.json"), encoding="utf-8") as f:
            return json.load(f).get("base_model_name_or_path") == MODEL_ID
    except Exception:
        return False


if os.path.isdir(LORA_DIR) and os.path.exists(os.path.join(LORA_DIR, "adapter_config.json")):
    if _lora_matches_base():
        from peft import PeftModel
        model = PeftModel.from_pretrained(model, LORA_DIR)
        model.eval()
        print(f"✅ تم تحميل طبقات LoRA المدرّبة من {LORA_DIR}")
    else:
        # LoRA مدرّبة على نموذج آخر (مثل Gemma) لا تصلح هنا — نؤرشفها ونبدأ من جديد
        import time as _t
        backup = f"{LORA_DIR}_old_{int(_t.time())}"
        os.rename(LORA_DIR, backup)
        print(f"⚠️ LoRA القديمة مبنية على نموذج مختلف — نُقلت إلى {backup}")
print("جاهز — النموذج محمّل.")

TRAIN_FLAG = os.environ.get("TRAIN", "0") == "1"
TRAIN_DATA = os.environ.get("TRAIN_DATA", "train.jsonl")
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "ddgoai-lora")
EPOCHS = int(os.environ.get("EPOCHS", "1"))
MAX_LEN = int(os.environ.get("MAX_LEN", "384"))            # أقصى طول للعينة بالتوكن
BATCH = int(os.environ.get("BATCH", "2"))                  # حجم الدفعة الصغيرة
GRAD_ACC = int(os.environ.get("GRAD_ACC", "8"))            # تجميع الخطوات (الدفعة الفعلية = BATCH×GRAD_ACC)
SAVE_EVERY = int(os.environ.get("SAVE_EVERY", "10"))       # حفظ LoRA كل N خطوة
MAX_TRAIN_SAMPLES = int(os.environ.get("MAX_TRAIN_SAMPLES", "0"))  # 0 = كل البيانات
LEARNING_RATE = float(os.environ.get("LEARNING_RATE", "2e-4"))

# ملف يتتبّع تقدّم التدريب عبر عدة جلسات متتالية (كل جلسة محدودة بمهلة
# ZeroGPU القصوى)، بحيث يُستكمَل التدريب تلقائياً من حيث توقّف بدل
# إعادته من الصفر في كل مرة.
TRAIN_STATE_FILE = os.path.join(OUTPUT_DIR, "training_state.json")
# هامش أمان تحت سقف ZeroGPU (300s) حتى تُحفَظ النتيجة قبل انتهاء المهلة.
CHUNK_TIME_BUDGET_SECONDS = int(os.environ.get("CHUNK_TIME_BUDGET_SECONDS", str(6 * 3600)))

import sys
from datetime import datetime

def log(msg):
    stamp = datetime.now().strftime("%H:%M:%S")
    print(f"[{stamp}] {msg}", flush=True)
    sys.stdout.flush()

EVAL_COUNT = 0


def load_train_dataset(path):
    """
    يقرأ ملف بيانات التدريب (JSON Lines) بطريقة متسامحة تضمن وصول كل
    سؤال وجواب موجود في الملف إلى التدريب 100%، بدل الاعتماد المباشر
    على pyarrow (عبر load_dataset) الذي يفشل بالكامل ويوقف التدريب
    كله لو سطر واحد فقط فيه عطل بسيط (مثل newline خام غير مهروب داخل
    نص، أو اقتباس غير مغلق).
    المنطق:
      1) نحاول تحليل كل سطر خام (مفصول بـ \n حقيقي) كـ JSON مستقل.
      2) أي سطر يفشل، نحاول إصلاحه بإلصاقه بالسطر التالي (لأن الفشل
         الشائع هو newline خام داخل نص لم يُهرَّب إلى \\n)، ونكرر
         المحاولة حتى ينجح التحليل أو حتى ننفد من الأسطر.
      3) لا نتجاهل أي سجل نهائياً إلا إذا تعذّر إصلاحه فعلياً
         (تالف بالكامل)، وفي هذه الحالة النادرة نطبع تحذيراً واضحاً
         بدل الفشل الصامت أو إيقاف التدريب بالكامل.
    """
    from datasets import Dataset

    with open(path, "r", encoding="utf-8") as f:
        raw_text = f.read()

    raw_lines = raw_text.split("\n")

    records = []
    unrecoverable = 0
    buf = ""
    buf_start_line = None

    for i, line in enumerate(raw_lines, 1):
        if buf == "":
            if line.strip() == "":
                continue
            buf = line
            buf_start_line = i
        else:
            # نعيد لصق الأسطر المكسورة كـ newline مهروب بدل خام،
            # لأن هذا هو سبب الكسر الأكثر شيوعاً في هذا النوع من الملفات.
            buf += "\\n" + line

        try:
            obj = json.loads(buf)
            # تحقق بسيط من الشكل المتوقع قبل القبول
            if isinstance(obj, dict) and "messages" in obj:
                records.append(obj)
            else:
                log(f"⚠️ سطر {buf_start_line}: صيغة JSON صحيحة لكن لا يحتوي على 'messages' — تم تجاوزه")
            buf = ""
            buf_start_line = None
        except json.JSONDecodeError:
            # لسا ناقص، نكمل نلصق بالسطر اللي بعده
            continue

    if buf != "":
        unrecoverable += 1
        log(f"⚠️ سطر {buf_start_line}: تعذّر إصلاحه تلقائياً (بقي غير مكتمل) — تم استبعاده من التدريب")
        log(f"   بداية المحتوى المتأثر: {buf[:150]}")

    total_raw_records = len(records) + unrecoverable
    log(f"📊 قراءة بيانات التدريب: {len(records)} سجل تم تحميله بنجاح من أصل {total_raw_records} "
        f"({'لا يوجد فقد — 100%' if unrecoverable == 0 else f'{unrecoverable} سجل تعذّر إصلاحه'})")

    return Dataset.from_list(records)


def read_training_state():
    if os.path.exists(TRAIN_STATE_FILE):
        try:
            with open(TRAIN_STATE_FILE, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception:
            pass
    return None


def write_training_state(state):
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    with open(TRAIN_STATE_FILE, "w", encoding="utf-8") as f:
        json.dump(state, f, ensure_ascii=False, indent=2)


# حالة التدريب الحية (تُقرأ من الواجهة كل بضع ثوانٍ)
TRAIN = {"running": False, "stop_requested": False, "msg": "لم يبدأ التدريب بعد.", "pct": 0.0}


def _fmt_time(sec):
    sec = int(max(sec, 0))
    h, r = divmod(sec, 3600)
    m, _ = divmod(r, 60)
    return f"{h}س {m}د" if h else f"{m}د"


def _content_to_text(c):
    if isinstance(c, str):
        return c
    if isinstance(c, list):
        return "".join(
            p.get("text", "") for p in c if isinstance(p, dict) and p.get("type") == "text"
        )
    return str(c)


def run_training_chunk():
    """
    يدرّب حتى اكتمال الخطوات المتبقية (أو حتى نفاد CHUNK_TIME_BUDGET_SECONDS)،
    ويحفظ LoRA وحالة التقدم كل SAVE_EVERY خطوة، فلا يضيع شيء لو توقف الـ Space.
    يُرجع (status, message, progress_fraction).
    """
    import time
    import math
    from peft import LoraConfig, get_peft_model, PeftModel
    from transformers import Trainer, TrainingArguments, DataCollatorForSeq2Seq, TrainerCallback

    global model
    t0 = time.time()

    state = read_training_state()
    existing_lora = os.path.isdir(OUTPUT_DIR) and os.path.exists(
        os.path.join(OUTPUT_DIR, "adapter_config.json")
    )

    # نغلّف النموذج بـ PEFT مرة واحدة فقط طوال عمر العملية
    if isinstance(model, PeftModel):
        peft_model = model
        for n_, p_ in peft_model.named_parameters():
            if "lora_" in n_:
                p_.requires_grad = True
    elif existing_lora:
        log(f"📂 تم العثور على LoRA سابقة في {OUTPUT_DIR} — سيتم التدريب فوقها (تراكمي)")
        peft_model = PeftModel.from_pretrained(model, OUTPUT_DIR, is_trainable=True)
    else:
        log(f"🆕 تدريب جديد من النموذج الأساسي {MODEL_ID}")
        lora_cfg = LoraConfig(
            r=16,
            lora_alpha=32,
            lora_dropout=0.05,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                            "gate_proj", "up_proj", "down_proj"],
            task_type="CAUSAL_LM",
        )
        peft_model = get_peft_model(model, lora_cfg)
    model = peft_model
    peft_model.print_trainable_parameters()

    raw = load_train_dataset(TRAIN_DATA)
    if MAX_TRAIN_SAMPLES > 0:
        raw = raw.shuffle(seed=42).select(range(min(MAX_TRAIN_SAMPLES, len(raw))))

    def tokenize(example):
        msgs = [{"role": m["role"], "content": _content_to_text(m["content"])}
                for m in example["messages"]]
        if msgs and msgs[-1]["role"] == "assistant":
            prompt_msgs, answer = msgs[:-1], msgs[-1]["content"]
        else:
            prompt_msgs, answer = msgs, ""
        template_kwargs = dict(tokenize=False, add_generation_prompt=True)
        if MODEL_SUPPORTS_NATIVE_THINKING:
            template_kwargs["enable_thinking"] = False
        prompt = tokenizer.apply_chat_template(prompt_msgs, **template_kwargs)
        full = prompt + answer + "<|im_end|>"
        p_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"]
        ids = tokenizer(full, add_special_tokens=False,
                        truncation=True, max_length=MAX_LEN)["input_ids"]
        k = min(len(p_ids), len(ids))
        # نحسب الخسارة على إجابة المساعد فقط (لا نتعلّم تكرار السؤال)
        labels = [-100] * k + ids[k:]
        return {"input_ids": ids, "attention_mask": [1] * len(ids), "labels": labels}

    train_ds = raw.map(tokenize, remove_columns=raw.column_names)
    train_ds = train_ds.filter(lambda e: any(l != -100 for l in e["labels"]))
    lens = train_ds["input_ids"]
    log(f"📏 {len(train_ds)} عينة صالحة | متوسط الطول {sum(map(len, lens)) / max(len(lens), 1):.0f} توكن")

    eff_batch = BATCH * GRAD_ACC
    steps_per_epoch = math.ceil(len(train_ds) / eff_batch)
    total_steps = steps_per_epoch * EPOCHS

    if state is None:
        state = {"total_steps": total_steps, "completed_steps": 0, "status": "in_progress"}
        write_training_state(state)
    if state.get("total_steps") != total_steps:
        state["total_steps"] = total_steps
        if state.get("completed_steps", 0) < total_steps and state.get("status") == "done":
            state["status"] = "in_progress"
        write_training_state(state)

    if state.get("status") == "done" and state.get("completed_steps", 0) >= total_steps:
        log("✅ التدريب مكتمل بالفعل.")
        return "done", "التدريب مكتمل بالفعل.", 1.0

    base_completed = state.get("completed_steps", 0)
    remaining_steps = max(total_steps - base_completed, 0)
    if remaining_steps <= 0:
        state["status"] = "done"
        write_training_state(state)
        return "done", "التدريب مكتمل بالفعل.", 1.0

    def save_progress(done_now, status="in_progress"):
        peft_model.save_pretrained(OUTPUT_DIR)
        state["completed_steps"] = done_now
        state["status"] = status
        write_training_state(state)

    class TimeBudgetCallback(TrainerCallback):
        def _chk(self, control):
            if time.time() - t0 > CHUNK_TIME_BUDGET_SECONDS:
                control.should_training_stop = True
            if TRAIN.get("stop_requested"):
                control.should_training_stop = True

        def on_substep_end(self, args, state_, control, **kwargs):
            self._chk(control)

        def on_step_end(self, args, state_, control, **kwargs):
            self._chk(control)

    class ChunkProgressCallback(TrainerCallback):
        def on_step_end(self, args, state_, control, **kwargs):
            done_now = base_completed + state_.global_step
            per = (time.time() - t0) / max(state_.global_step, 1)
            eta = per * (total_steps - done_now)
            pct = 100.0 * done_now / max(total_steps, 1)
            TRAIN["pct"] = round(pct, 1)
            TRAIN["msg"] = (f"⏳ خطوة {done_now}/{total_steps} ({pct:.1f}%) — "
                            f"{per:.0f} ثانية/خطوة — المتبقي ≈ {_fmt_time(eta)}")
            if state_.global_step % SAVE_EVERY == 0:
                save_progress(done_now)

        def on_log(self, args, state_, control, logs=None, **kwargs):
            if logs and "loss" in logs:
                done_now = base_completed + state_.global_step
                log(f"خطوة {done_now}/{total_steps} | الخسارة: {logs['loss']:.4f}")

    targs = dict(
        output_dir=OUTPUT_DIR,
        max_steps=remaining_steps,
        per_device_train_batch_size=BATCH,
        gradient_accumulation_steps=GRAD_ACC,
        learning_rate=LEARNING_RATE,
        lr_scheduler_type="constant_with_warmup",
        warmup_steps=5,
        logging_steps=1,
        save_strategy="no",
        eval_strategy="no",
        report_to="none",
        use_cpu=True,
        seed=42 + base_completed,       # ترتيب مختلف لكل جزء، فلا نكرر نفس العينات
        dataloader_num_workers=0,
    )
    try:
        args = TrainingArguments(group_by_length=True, **targs)   # يقلّل الحشو (padding)
    except TypeError:
        args = TrainingArguments(**targs)

    trainer = Trainer(
        model=peft_model,
        args=args,
        train_dataset=train_ds,
        data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True, label_pad_token_id=-100),
        callbacks=[TimeBudgetCallback(), ChunkProgressCallback()],
    )

    log(f"▶️ بدء التدريب — المنجز سابقاً: {base_completed}/{total_steps} | "
        f"المتبقي: {remaining_steps} خطوة (دفعة فعلية {eff_batch})")
    base_cfg = peft_model.get_base_model().config
    base_cfg.use_cache = False
    try:
        trainer.train()
    finally:
        base_cfg.use_cache = True
        peft_model.eval()

    new_completed = min(base_completed + trainer.state.global_step, total_steps)
    tokenizer.save_pretrained(OUTPUT_DIR)

    if new_completed >= total_steps:
        save_progress(new_completed, "done")
        log(f"🏁 اكتمل التدريب: {new_completed}/{total_steps}. الحفظ في {OUTPUT_DIR}")
        return "done", f"✅ اكتمل التدريب ({new_completed}/{total_steps} خطوة).", 1.0

    save_progress(new_completed, "in_progress")
    pct = 100.0 * new_completed / total_steps
    if TRAIN.get("stop_requested"):
        log(f"⏸️ أُوقف التدريب يدوياً: {new_completed}/{total_steps} ({pct:.1f}%)")
        return "stopped", f"⏸️ تم الإيقاف عند {new_completed} من {total_steps} خطوة ({pct:.1f}%). النموذج الحالي محمّل وجاهز للتجربة في الدردشة.", new_completed / total_steps
    log(f"💾 حُفظ تقدّم جزئي: {new_completed}/{total_steps} ({pct:.1f}%)")
    return "in_progress", f"تم إنجاز {new_completed} من {total_steps} خطوة ({pct:.1f}%).", new_completed / total_steps


def _train_worker(auto_continue):
    """يعمل في خيط خلفي: يستمر التدريب حتى لو أُغلقت الصفحة أو انقطع الاتصال،
    إلى أن يكتمل أو يُطلب إيقافه من زر الإيقاف."""
    try:
        status = "in_progress"
        while status == "in_progress":
            status, message, frac = run_training_chunk()
            TRAIN["pct"] = round(frac * 100, 1)
            TRAIN["msg"] = message
            if status == "stopped" or not auto_continue:
                break
    except Exception as e:
        traceback.print_exc()
        TRAIN["msg"] = f"❌ {type(e).__name__}: {e}"
    finally:
        TRAIN["running"] = False
        TRAIN["stop_requested"] = False


def poll_training():
    if not TRAIN["running"] and TRAIN["msg"].startswith("لم يبدأ"):
        st = read_training_state()
        if st and st.get("total_steps"):
            pct = round(100 * st.get("completed_steps", 0) / st["total_steps"], 1)
            return f"محفوظ سابقاً: {st.get('completed_steps', 0)}/{st['total_steps']} خطوة ({pct}%).", pct
    return TRAIN["msg"], TRAIN["pct"]


_SELF_TALK_RULE = (
    "محتوى تفكيرك يجب أن يكون حديثاً داخلياً مع نفسك فقط، كأنك تفكر بصوت خافت "
    "قبل أن يسمعك أحد. استخدم ضمير المتكلم (أنا، دعني أرى، هل هذا صحيح؟) "
    "واطرح على نفسك أسئلة وراجع إجاباتك عليها. ممنوع منعاً باتاً أن تخاطب "
    "المستخدم داخل <think> أو تستخدم أي كلمة موجهة له مثل 'أنت' أو 'يمكنك' "
    "أو 'إذا كنت تقصد' أو 'سأجيبك' — هذا الكلام يُكتب فقط بعد إغلاق </think>."
)

THINK_MIN_TOKENS = {
    "non": 0,
    "low": 150,
    "medium": 500,
    "high": 1200,
    "xhigh": 2000,
    "max": 3000,
}

EFFORT_CONFIG = {
    "non": dict(answer_tokens=MAX_OUTPUT_TOKENS, instruction=""),
    "low": dict(
        answer_tokens=MAX_OUTPUT_TOKENS,
        instruction=(
            f"{_SELF_TALK_RULE} قبل إجابتك، فكّر مع نفسك داخل <think> و </think> "
            "بإيجاز شديد جداً، فقط 3 أو 4 أسطر قصيرة كأنك تسأل نفسك سؤالاً "
            "وتجيب عليه بسرعة، ثم أغلق </think> واكتب الإجابة الموجهة للمستخدم."
        ),
    ),
    "medium": dict(
        answer_tokens=MAX_OUTPUT_TOKENS,
        instruction=(
            f"{_SELF_TALK_RULE} قبل إجابتك، فكّر مع نفسك داخل <think> و </think> "
            "في حدود 10 إلى 12 سطراً، كحوار داخلي متوازن تسأل فيه نفسك وتراجع "
            "استنتاجك، ثم أغلق </think> واكتب الإجابة الموجهة للمستخدم."
        ),
    ),
    "high": dict(
        answer_tokens=MAX_OUTPUT_TOKENS,
        instruction=(
            f"{_SELF_TALK_RULE} قبل إجابتك، فكّر مع نفسك داخل <think> و </think> "
            "بعمق وتأنٍّ، اطرح على نفسك أسئلة من عدة زوايا وراجع كل احتمال قبل "
            "أن تستقر على استنتاج، لا تستعجل في إنهاء حوارك الداخلي، ثم أغلق "
            "</think> واكتب الإجابة الموجهة للمستخدم."
        ),
    ),
    "xhigh": dict(
        answer_tokens=MAX_OUTPUT_TOKENS,
        instruction=(
            f"{_SELF_TALK_RULE} قبل إجابتك، فكّر مع نفسك داخل <think> و </think> "
            "بعمق كبير جداً ولفترة طويلة جداً، اسأل نفسك أسئلة متتالية وحلل "
            "المشكلة من كل الزوايا الممكنة، تحقق من أخطائك المحتملة أكثر من "
            "مرة بمراجعة ذاتية صريحة، لا تستعجل في إنهاء حوارك الداخلي إطلاقاً، "
            "ثم أغلق </think> واكتب الإجابة الموجهة للمستخدم."
        ),
    ),
    "max": dict(
        answer_tokens=MAX_OUTPUT_TOKENS,
        instruction=(
            f"{_SELF_TALK_RULE} قبل إجابتك، فكّر مع نفسك داخل <think> و </think> "
            "بأقصى عمق ممكن ولفترة طويلة جداً حتى لو كان السؤال بسيطاً جداً. "
            "اسأل نفسك أكثر من سؤال، حلل كل جانب بالتفصيل كمونولوج داخلي خالص، "
            "فكر في أكثر من صياغة للإجابة وقارن بينها مع نفسك، حسّنها، تحقق من "
            "دقتها مرات عديدة بمراجعة ذاتية، استمر في حوارك الداخلي لأطول فترة "
            "ممكنة قبل كتابة الإجابة الموجهة للمستخدم خارج </think>."
        ),
    ),
}

def build_messages(message, history, images, cfg):
    messages = []

    if cfg["instruction"]:
        messages.append({
            "role": "system",
            "content": (
                "التزم حرفياً بالتعليمة التالية في أسلوب كتابة ردك: "
                f"{cfg['instruction']} "
                "اكتب قسم التفكير فعلياً داخل وسمي <think> و </think> في بداية كل "
                "رد قبل الإجابة النهائية. هذا مجرد أسلوب كتابة مطلوب منك، وليس له "
                "أي علاقة بإفشاء تعليمات سرية أو حساسة."
            ),
        })

    for turn in history:
        # turn["content"] قد يصل أحياناً كقائمة (list) بدل نص عادي، خاصة
        # عندما يحتوي تاريخ Gradio على صور أو عناصر مرفقة ضمن نفس الرسالة.
        # قالب Qwen3 كان يتسامح مع هذا، لكن قالب Qwen2.5-Coder يتوقع نصاً
        # صرفاً (str) فقط ويفشل بخطأ "can only concatenate str (not list)"
        # إن لم نحوّله أولاً. _content_to_text تتعامل مع الحالتين بأمان.
        messages.append({"role": turn["role"], "content": _content_to_text(turn["content"])})

    messages.append({"role": "user", "content": _content_to_text(message)})
    return messages

from transformers import LogitsProcessor, LogitsProcessorList

class MinThinkProcessor(LogitsProcessor):
    def __init__(self, think_end_ids, prompt_len, min_tokens):
        self.think_end_ids = think_end_ids
        self.prompt_len = prompt_len
        self.min_tokens = min_tokens

    def __call__(self, input_ids, scores):
        generated = input_ids.shape[1] - self.prompt_len
        if generated < self.min_tokens and self.think_end_ids:
            first = self.think_end_ids[0]
            scores[:, first] = -float("inf")
        return scores

def generate_response(message, history, effort, images):
    cfg = EFFORT_CONFIG[effort]
    messages = build_messages(message, history, images, cfg)

    template_kwargs = dict(tokenize=False, add_generation_prompt=True)
    if MODEL_SUPPORTS_NATIVE_THINKING:
        # عند Qwen3 فقط: نستخدم علم القالب الأصلي أيضاً (طبقة حماية إضافية
        # فوق حقن <think>\n اليدوي أدناه، لا تعارض بينهما).
        template_kwargs["enable_thinking"] = bool(cfg["instruction"])
    text_prompt = tokenizer.apply_chat_template(messages, **template_kwargs)

    # حقن <think>\n يدوياً يعمل بنفس الطريقة لأي نموذج (Qwen3 أو
    # Qwen2.5-Coder أو غيرهما)، لأن التفكير هنا أسلوب كتابة نتعلّمه عبر
    # system prompt + forced prefix، وليس ميزة مدمجة بالنموذج تحديداً.
    forced_prefix = ""
    if cfg["instruction"]:
        forced_prefix = "<think>\n"
        text_prompt = text_prompt + forced_prefix

    inputs = tokenizer(text_prompt, return_tensors="pt", add_special_tokens=False)

    input_len = inputs["input_ids"].shape[1]
    if input_len >= MAX_CONTEXT_TOKENS:
        yield (
            "⚠️ انتهت مساحة السياق المتاحة لهذا النموذج على هذه الموارد "
            f"({MAX_CONTEXT_TOKENS} توكن). الرجاء مسح المحادثة والبدء من جديد."
        )
        return

    allowed_new_tokens = min(cfg["answer_tokens"], MAX_CONTEXT_TOKENS - input_len)

    streamer = TextIteratorStreamer(
        tokenizer,
        skip_prompt=True,
        skip_special_tokens=False
    )

    logits_processors = LogitsProcessorList()
    if forced_prefix and THINK_MIN_TOKENS.get(effort, 0) > 0:
        think_end_ids = tokenizer.encode("</think>", add_special_tokens=False)
        logits_processors.append(
            MinThinkProcessor(think_end_ids, input_len, THINK_MIN_TOKENS[effort])
        )

    gen_kwargs = dict(
        **inputs,
        max_new_tokens=allowed_new_tokens,
        do_sample=True,
        temperature=0.6 if forced_prefix else 0.7,
        top_p=0.95 if forced_prefix else 0.8,
        top_k=20,
        streamer=streamer,
        logits_processor=logits_processors,
    )
    thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()

    full = forced_prefix
    for chunk in streamer:
        full += chunk
        yield full
    thread.join()

CLEAN_PATTERNS = re.compile(r"<\|?turn\|?>|<\|?channel\|?>|<\|?end\|?>|<eos>|<bos>|<\|im_start\|>|<\|im_end\|>|<\|endoftext\|>")

def clean_output(text):
    return CLEAN_PATTERNS.sub("", text)

def split_think_answer(full_text):
    full_text = clean_output(full_text)
    if "<think>" in full_text and "</think>" in full_text:
        think = full_text.split("<think>", 1)[1].split("</think>", 1)[0].strip()
        answer = full_text.split("</think>", 1)[1].strip()
        return think, answer
    elif "<think>" in full_text:
        think = full_text.split("<think>", 1)[1].strip()
        return think, ""
    else:
        return "", full_text.strip()

def render_message(think, answer, thinking_done):
    if not think:
        return answer
    status = "اكتمل التفكير" if thinking_done else "قاعد يفكر..."
    html = f"""<details {"open" if not thinking_done else ""} class="think-box">
<summary>🧠 {status}</summary>
<div class="think-content">{think}</div>
</details>
{answer}"""
    return html

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Tajawal:wght@400;500;700;900&family=JetBrains+Mono:wght@400;500&family=Lora:wght@500;600&display=swap');
:root {
    --ink: #0d0d0d;
    --panel: #1a1a1a;
    --panel-2: #262624;
    --line: #3a3a38;
    --cyan: #d97757;
    --cyan-dim: #b35f44;
    --text: #e8e6e1;
    --text-dim: #8e8c87;
}
* { box-sizing: border-box; }
.gradio-container {
    background: var(--ink) !important;
    font-family: 'Tajawal', sans-serif !important;
    color: var(--text) !important;
    max-width: 880px !important;
    margin: 0 auto !important;
}
footer { display: none !important; }
#header {
    text-align: center;
    padding: 90px 20px 60px 20px;
    border-bottom: none;
    margin-bottom: 0;
}
#header .spark {
    font-size: 2.6rem;
    color: var(--cyan);
    margin-bottom: 22px;
    display: inline-block;
    animation: spark-spin 10s linear infinite;
}
@keyframes spark-spin {
    from { transform: rotate(0deg); }
    to { transform: rotate(360deg); }
}
#header h1 {
    font-family: 'Lora', serif;
    font-weight: 600;
    font-size: 2.3rem;
    letter-spacing: -0.01em;
    margin: 0 0 6px 0;
    color: var(--text);
}
#header .tag {
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.72rem;
    color: var(--text-dim);
    letter-spacing: 0.06em;
    opacity: 0.7;
}
#header .desc {
    color: var(--text-dim);
    font-size: 0.95rem;
    margin-top: 14px;
    line-height: 1.6;
}
.chatbot {
    background: transparent !important;
    border: none !important;
}
.message-wrap { direction: rtl; }
.message.user {
    background: var(--panel-2) !important;
    color: var(--text) !important;
    border-radius: 18px 18px 4px 18px !important;
}
.message.bot {
    background: transparent !important;
    color: var(--text) !important;
    border: none !important;
    border-radius: 0 !important;
}
.think-box {
    margin-bottom: 10px;
    border: 1px solid var(--line);
    border-radius: 10px;
    background: #0e1319;
    overflow: hidden;
}
.think-box summary {
    cursor: pointer;
    padding: 8px 12px;
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.75rem;
    color: var(--cyan);
    list-style: none;
    user-select: none;
}
.think-box summary::-webkit-details-marker { display: none; }
.think-content {
    padding: 10px 14px 14px 14px;
    color: var(--text-dim);
    font-size: 0.85rem;
    line-height: 1.7;
    white-space: pre-wrap;
    border-top: 1px solid var(--line);
}
#input-row {
    background: var(--panel-2) !important;
    border: 1px solid var(--line) !important;
    border-radius: 28px !important;
    padding: 10px 14px !important;
    align-items: center !important;
}
#input-row textarea {
    background: transparent !important;
    border: none !important;
    color: var(--text) !important;
    font-family: 'Tajawal', sans-serif !important;
    font-size: 1rem !important;
    box-shadow: none !important;
}
#input-row textarea:focus {
    border: none !important;
    box-shadow: none !important;
}
button.primary {
    background: var(--cyan) !important;
    color: #1a1a1a !important;
    border: none !important;
    font-weight: 700 !important;
    border-radius: 50% !important;
    aspect-ratio: 1 / 1 !important;
    min-width: 42px !important;
    width: 42px !important;
    padding: 0 !important;
}
button.primary:hover { background: #e38a68 !important; }
#stop-btn {
    background: #e5484d !important;
    color: #fff !important;
    border: none !important;
    border-radius: 50% !important;
    aspect-ratio: 1 / 1 !important;
    font-size: 0.9rem !important;
    font-weight: 700 !important;
    min-width: 42px !important;
    width: 42px !important;
    padding: 0 !important;
}
#stop-btn:hover {
    background: #c93a3e !important;
}
button.secondary {
    background: transparent !important;
    border: 1px solid var(--line) !important;
    color: var(--text-dim) !important;
    border-radius: 12px !important;
}
button.secondary:hover {
    border-color: var(--cyan) !important;
    color: var(--cyan) !important;
}
#effort-radio label {
    font-family: 'JetBrains Mono', monospace !important;
    font-size: 0.8rem !important;
}
#footer-note {
    text-align: center;
    color: var(--text-dim);
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.72rem;
    padding: 18px 0 30px 0;
    letter-spacing: 0.03em;
}
"""

EFFORT_LABELS = {
    "non": "Non — بدون تفكير",
    "low": "Low — تفكير قليل",
    "medium": "Medium — تفكير عادي",
    "high": "High — تفكير طويل",
    "xhigh": "XHigh — تفكير مضاعف",
    "max": "Max — أقصى تحليل",
}
LABEL_TO_KEY = {v: k for k, v in EFFORT_LABELS.items()}

with gr.Blocks(title="Ddgoai") as demo:
    gr.HTML(
        """
        <div id="header">
            <div class="spark">✦</div>
            <h1>Ddgoai</h1>
            <div class="desc">اكتب سؤالك للبدء</div>
        </div>
        """
    )

    chatbot = gr.Chatbot(
        elem_classes="chatbot",
        height=480,
        show_label=False,
        avatar_images=(None, None),
        rtl=True,
    )

    effort = gr.Radio(
        choices=list(EFFORT_LABELS.values()),
        value=EFFORT_LABELS["non"],
        label="Think Effort",
        elem_id="effort-radio",
    )

    images_box = gr.Gallery(
        label="الصور المرفقة (اختياري)",
        show_label=True,
        columns=4,
        height=120,
        visible=False,
    )

    with gr.Row(elem_id="input-row"):
        msg = gr.Textbox(
            placeholder="اكتب رسالتك هنا...",
            show_label=False,
            scale=6,
            rtl=True,
            container=False,
        )
        img_upload = gr.File(
            label="",
            file_count="multiple",
            file_types=["image"],
            scale=2,
            visible=False,
        )
        send = gr.Button("إرسال", variant="primary", scale=1, elem_id="send-btn")
        stop = gr.Button("■", variant="stop", scale=1, elem_id="stop-btn", visible=False)

    with gr.Row():
        clear = gr.Button("مسح المحادثة", variant="secondary", size="sm")

    with gr.Accordion("🎓 تدريب النموذج", open=False):
        gr.Markdown(
            "التدريب يعمل في الخلفية على CPU ويحفظ تقدمه كل عدة خطوات."
            " يمكنك إغلاق الصفحة والعودة لاحقاً، وسيستكمل من حيث توقف."
            " الدردشة متوقفة أثناء التدريب."
        )
        train_status = gr.Markdown("لم يبدأ التدريب بعد.")
        train_progress = gr.Slider(
            label="التقدّم", minimum=0, maximum=100, value=0,
            interactive=False, show_label=True
        )
        with gr.Row():
            train_btn = gr.Button("🚀 ابدأ/أكمل التدريب", variant="primary")
            train_stop_btn = gr.Button("⏸️ إيقاف التدريب", variant="stop")
            train_auto = gr.Checkbox(label="استكمال تلقائي حتى الاكتمال", value=True)
        gr.Markdown(
            "بعد الإيقاف يمكنك تجربة آخر نسخة مدرّبة مباشرة في الدردشة بالأعلى، "
            "ثم الضغط على *ابدأ/أكمل التدريب* لاحقاً لمتابعة نفس التقدّم."
        )

    gr.HTML('<div id="footer-note">THE ZYZ STUDIO — Ddgoai (Qwen2.5-Coder-1.5B)</div>')

    SHOW_SEND = (gr.update(visible=True), gr.update(visible=False))
    SHOW_STOP = (gr.update(visible=False), gr.update(visible=True))

    def respond(message, history, effort_label, uploaded_files):
        if TRAIN["running"]:
            history = history + [
                {"role": "user", "content": message or ""},
                {"role": "assistant", "content": "⏳ التدريب جارٍ الآن، الدردشة متوقفة مؤقتاً لتسريعه."},
            ]
            yield history, "", None, [], gr.update(visible=True), gr.update(visible=False)
            return
        if (not message or not message.strip()) and not uploaded_files:
            send_u, stop_u = SHOW_SEND
            yield history, "", None, [], send_u, stop_u
            return

        from PIL import Image
        images = []
        file_names = []
        if uploaded_files:
            for f in uploaded_files:
                file_names.append(os.path.basename(f.name))
                try:
                    images.append(Image.open(f.name).convert("RGB"))
                except Exception:
                    pass

        effort_key = LABEL_TO_KEY[effort_label]
        base_text = message if message else "(صورة بدون نص)"
        if file_names:
            files_label = "📎 الملفات المرفقة:\n" + "\n".join(f"• {name}" for name in file_names)
            display_text = f"{files_label}\n\n{base_text}"
        else:
            display_text = base_text

        history = history + [{"role": "user", "content": display_text}]
        history = history + [{"role": "assistant", "content": ""}]

        send_u, stop_u = SHOW_STOP
        yield history, "", None, [], send_u, stop_u

        import time
        prior_history = history[:-2]
        shown = ""

        for partial in generate_response(message or "صف هذه الصورة", prior_history, effort_key, images):
            new_part = partial[len(shown):]
            for ch in new_part:
                shown += ch
                think, answer = split_think_answer(shown)
                thinking_done = "</think>" in shown or not think
                rendered = render_message(think, answer, thinking_done)
                history[-1] = {"role": "assistant", "content": rendered}
                yield history, "", None, [], send_u, stop_u
                time.sleep(0.008)

        think, answer = split_think_answer(shown)
        rendered = render_message(think, answer, True)
        history[-1] = {"role": "assistant", "content": rendered}
        send_u, stop_u = SHOW_SEND
        yield history, "", None, [], send_u, stop_u

    submit_event_1 = msg.submit(
        respond,
        [msg, chatbot, effort, img_upload],
        [chatbot, msg, img_upload, images_box, send, stop]
    )

    submit_event_2 = send.click(
        respond,
        [msg, chatbot, effort, img_upload],
        [chatbot, msg, img_upload, images_box, send, stop]
    )

    def on_stop():
        send_u, stop_u = SHOW_SEND
        return send_u, stop_u

    stop.click(
        on_stop,
        None,
        [send, stop],
        cancels=[submit_event_1, submit_event_2],
    )

    clear.click(
        lambda: ([], "", None, [], *SHOW_SEND),
        None,
        [chatbot, msg, img_upload, images_box, send, stop],
        cancels=[submit_event_1, submit_event_2],
    )

    def start_training(auto_continue):
        if TRAIN["running"]:
            return TRAIN["msg"], TRAIN["pct"]
        TRAIN["running"] = True
        TRAIN["msg"] = "⏳ جارٍ التحضير (قراءة البيانات وتجهيزها)..."
        threading.Thread(target=_train_worker, args=(bool(auto_continue),), daemon=True).start()
        return TRAIN["msg"], TRAIN["pct"]

    def request_stop():
        if TRAIN["running"]:
            TRAIN["stop_requested"] = True
            return "⏸️ جارٍ الإيقاف بعد نهاية الخطوة الحالية...", TRAIN["pct"]
        return TRAIN["msg"], TRAIN["pct"]

    train_btn.click(start_training, [train_auto], [train_status, train_progress])
    train_stop_btn.click(request_stop, None, [train_status, train_progress])

    train_timer = gr.Timer(3)
    train_timer.tick(poll_training, None, [train_status, train_progress])
    demo.load(poll_training, None, [train_status, train_progress])

demo.launch(server_name="0.0.0.0", server_port=7860, css=CSS, theme=gr.themes.Base())