Blame for app.py

import os, io, re, gc, json, time, base64, threading
from fnmatch import fnmatch
import numpy as np
import torch
import soundfile as sf
from PIL import Image
from gradio import Server
from fastapi.responses import HTMLResponse
from huggingface_hub import snapshot_download, HfApi, constants
from transformers import (AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM,
                          WhisperProcessor, WhisperForConditionalGeneration,
                          VisionEncoderDecoderModel, ViTImageProcessor,
                          TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList)

try:  # ZeroGPU على Hugging Face (اختياري)
    import spaces
    GPU = spaces.GPU
except Exception:
    def GPU(duration=60):
        return lambda f: f

app = Server()
HERE = os.path.dirname(os.path.abspath(__file__))

# ---- إعدادات الستريمينج: نجمع التوكنات ثم نرسلها دفعة واحدة ليبدو الرد سلساً وسريعاً ----
BATCH_TOKENS = int(os.environ.get("BATCH_TOKENS", 6))     # عدد القطع قبل كل إرسال
BATCH_SECONDS = float(os.environ.get("BATCH_SECONDS", 0.25))  # أو أقصى انتظار بين إرسالين
MAX_HISTORY = 12                                          # آخر N رسالة تُرسل للنموذج


@GPU(duration=10)
def _gpu_marker():  # ZeroGPU يشترط وجود دالة واحدة على الأقل؛ النماذج تعمل على CPU
    return 0


def M(repo, kind, size, group, bf16=False):
    return {"repo": repo, "kind": kind, "size": size, "group": group, "bf16": bf16}

SM, MID, BIG = "💬 دردشة صغيرة (أقل من 1B)", "🧠 متوسطة (1–2B)", "🚀 كبيرة (2–4B) — بطيئة على CPU"
CODE, T5, CP, CU = "👨‍💻 برمجة", "⚡ سريعة (سؤال وجواب)", "🐣 إكمال نصي", "🧩 مخصّص"
MODELS = {
    # --- صغيرة ---
    "qwen25":    M("Qwen/Qwen2.5-0.5B-Instruct", "chat", "~1GB", SM),
    "qwen3":     M("Qwen/Qwen3-0.6B", "chat", "~1.2GB", SM),
    "qwen2":     M("Qwen/Qwen2-0.5B-Instruct", "chat", "~1GB", SM),
    "qwen":      M("Qwen/Qwen1.5-0.5B-Chat", "chat", "~1GB", SM),
    "lfm2":      M("LiquidAI/LFM2-350M", "chat", "~700MB", SM),
    "lfm2_700":  M("LiquidAI/LFM2-700M", "chat", "~1.4GB", SM),
    "smol360":   M("HuggingFaceTB/SmolLM2-360M-Instruct", "chat", "~720MB", SM),
    "smol135":   M("HuggingFaceTB/SmolLM2-135M-Instruct", "chat", "~270MB", SM),
    "qwen15_18": M("Qwen/Qwen1.5-1.8B-Chat", "chat", "~3.7GB", MID),
    # --- متوسطة ---
    "qwen25_15": M("Qwen/Qwen2.5-1.5B-Instruct", "chat", "~3GB", MID),
    "qwen3_17":  M("Qwen/Qwen3-1.7B", "chat", "~3.4GB", MID),
    "smol17":    M("HuggingFaceTB/SmolLM2-1.7B-Instruct", "chat", "~3.4GB", MID),
    "lfm2_12":   M("LiquidAI/LFM2-1.2B", "chat", "~2.4GB", MID),
    "tinyllama": M("TinyLlama/TinyLlama-1.1B-Chat-v1.0", "chat", "~2.2GB", MID),
    "r1_15":     M("deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B", "chat", "~3.5GB", MID),
    "granite2":  M("ibm-granite/granite-3.3-2b-instruct", "chat", "~5GB", MID, True),
    # --- كبيرة (bfloat16 لتوفير الذاكرة) ---
    "qwen25_3":  M("Qwen/Qwen2.5-3B-Instruct", "chat", "~6GB", BIG, True),
    "smol3":     M("HuggingFaceTB/SmolLM3-3B", "chat", "~6GB", BIG, True),
    "phi35":     M("microsoft/Phi-3.5-mini-instruct", "chat", "~7.6GB", BIG, True),
    "phi4mini":  M("microsoft/Phi-4-mini-instruct", "chat", "~7.6GB", BIG, True),
    "qwen3_4":   M("Qwen/Qwen3-4B-Instruct-2507", "chat", "~8GB", BIG, True),
    "lfm2_26":   M("LiquidAI/LFM2-2.6B", "chat", "~5.2GB", BIG, True),
    # --- برمجة ---
    "qcoder05":  M("Qwen/Qwen2.5-Coder-0.5B-Instruct", "chat", "~1GB", CODE),
    "qcoder15":  M("Qwen/Qwen2.5-Coder-1.5B-Instruct", "chat", "~3GB", CODE),
    "qcoder3":   M("Qwen/Qwen2.5-Coder-3B-Instruct", "chat", "~6GB", CODE, True),
    # --- سؤال وجواب ---
    "lamini":    M("MBZUAI/LaMini-Flan-T5-77M", "t5", "~300MB", T5),
    "lamini248": M("MBZUAI/LaMini-Flan-T5-248M", "t5", "~1GB", T5),
    "lamini783": M("MBZUAI/LaMini-Flan-T5-783M", "t5", "~3GB", T5),
    "flansmall": M("google/flan-t5-small", "t5", "~300MB", T5),
    "flanbase":  M("google/flan-t5-base", "t5", "~1GB", T5),
    "flanlarge": M("google/flan-t5-large", "t5", "~3GB", T5),
    "t5small":   M("google/t5-small", "t5", "~240MB", T5),
    # --- إكمال نصي ---
    "distilgpt2": M("distilbert/distilgpt2", "alpaca", "~350MB", CP),
    "gpt2":      M("openai-community/gpt2", "alpaca", "~500MB", CP),
    "opt125":    M("facebook/opt-125m", "alpaca", "~250MB", CP),
    "opt350":    M("facebook/opt-350m", "alpaca", "~650MB", CP),
    "lgpt":      M("MBZUAI/LaMini-GPT-124M", "alpaca", "~250MB", CP),
    "lgpt774":   M("MBZUAI/LaMini-GPT-774M", "alpaca", "~1.5GB", CP),
    # --- صوت وصور ---
    "whisper":   M("openai/whisper-tiny.en", "asr", "~150MB", "aux"),
    "caption":   M("nlpconnect/vit-gpt2-image-captioning", "caption", "~950MB", "aux"),
}
TEXT_KINDS = ("chat", "t5", "alpaca")

_pipes, _state, _cancels = {}, {}, {}
_lock = threading.Lock()
F32 = torch.float32


def _st(mid):
    return _state.setdefault(mid, {"state": "idle", "pct": 0, "error": "", "t0": 0, "took": 0,
                                   "done": 0, "total": 0, "speed": 0, "eta": 0})


def _plan(repo):
    """الملفات المطلوبة وحجمها الإجمالي (لحساب النسبة)."""
    sib = HfApi().model_info(repo, files_metadata=True).siblings
    ign = ["*.h5", "*.msgpack", "*.ot", "*.onnx", "*.onnx_data", "*.tflite",
           "*.gguf", "onnx/*", "openvino/*", "coreml/*", "original/*", "consolidated*"]
    if any(s.rfilename.endswith(".safetensors") for s in sib):
        ign.append("*.bin")
    keep = [s for s in sib if not any(fnmatch(s.rfilename, p) for p in ign)]
    return ign, sum(s.size or 0 for s in keep)


def _watch(mid, repo, total, stop):
    """يقيس حجم ما نزل في الكاش: النسبة + السرعة + الوقت المتبقي."""
    d = os.path.join(constants.HF_HUB_CACHE, "models--" + repo.replace("/", "--"))
    s = _st(mid)
    t_start, n0 = time.time(), None
    while not stop.is_set():
        n = 0
        for r, _, fs in os.walk(d):
            for f in fs:
                p = os.path.join(r, f)
                try:
                    if not os.path.islink(p):
                        n += os.path.getsize(p)
                except OSError:
                    pass
        if n0 is None:
            n0 = n  # ما كان موجوداً مسبقاً في الكاش لا يدخل في حساب السرعة
        dt = time.time() - t_start
        speed = (n - n0) / dt if dt > 1 else 0
        s["done"], s["total"], s["speed"] = n, total, speed
        s["eta"] = max(0, (total - n) / speed) if speed > 0 and total > n else 0
        if total:
            s["pct"] = min(99.0, n * 100 / total)
        stop.wait(0.7)


def _build(info):
    k, r = info["kind"], info["repo"]
    if k == "asr":
        return {"proc": WhisperProcessor.from_pretrained(r),
                "model": WhisperForConditionalGeneration.from_pretrained(r, dtype=F32).eval()}
    if k == "caption":
        return {"proc": ViTImageProcessor.from_pretrained(r), "tok": AutoTokenizer.from_pretrained(r),
                "model": VisionEncoderDecoderModel.from_pretrained(r, dtype=F32).eval()}
    cls = AutoModelForSeq2SeqLM if k == "t5" else AutoModelForCausalLM
    dt = torch.bfloat16 if info.get("bf16") else F32
    return {"tok": AutoTokenizer.from_pretrained(r), "model": cls.from_pretrained(r, dtype=dt).eval()}


def _free_text(keep):
    """نموذج نصي واحد في الذاكرة فقط، لتجنّب نفاد RAM مع النماذج الكبيرة."""
    for k in list(_pipes):
        if k != keep and MODELS.get(k, {}).get("kind") in TEXT_KINDS:
            _pipes.pop(k, None)
            _st(k).update(state="idle", pct=0)
    gc.collect()


def _load_sync(mid):
    info, s, stop = MODELS[mid], _st(mid), threading.Event()
    s.update(state="downloading", pct=0, error="", t0=time.time(), took=0,
             done=0, total=0, speed=0, eta=0)
    try:
        ign, total = _plan(info["repo"])
        s["total"] = total
        threading.Thread(target=_watch, args=(mid, info["repo"], total, stop), daemon=True).start()
        snapshot_download(info["repo"], ignore_patterns=ign)
        stop.set()
        s.update(state="loading", pct=100, eta=0)
        if info["kind"] in TEXT_KINDS:
            _free_text(mid)
        _pipes[mid] = _build(info)
        s.update(state="ready", took=time.time() - s["t0"])
    except Exception as e:
        stop.set()
        s.update(state="error", error=str(e)[:300])


class StopOnCancel(StoppingCriteria):
    def __init__(self, rid):
        self.rid = rid

    def __call__(self, input_ids, scores, **kw):
        return torch.full((input_ids.shape[0],), bool(_cancels.get(self.rid)),
                          dtype=torch.bool, device=input_ids.device)


def _alpaca(msgs):
    q = next((m["content"] for m in reversed(msgs) if m["role"] == "user"), "")
    return ("Below is an instruction that describes a task. Write a response that "
            f"appropriately completes the request.\n\n### Instruction:\n{q}\n\n### Response:\n")


# ============ API ============
@app.api(name="models")
def api_models() -> dict:
    return {"models": [{"id": k, **v} for k, v in MODELS.items()]}


@app.api(name="register")
def api_register(repo: str, kind: str = "chat") -> dict:
    repo = repo.strip()
    if not re.fullmatch(r"[\w.\-]+/[\w.\-]+", repo) or kind not in TEXT_KINDS:
        return {"ok": False, "error": "صيغة غير صحيحة. مثال: Qwen/Qwen2.5-3B-Instruct"}
    mid = "c_" + re.sub(r"\W", "_", repo)
    MODELS.setdefault(mid, M(repo, kind, "؟", CU, True))
    return {"ok": True, "model": {"id": mid, **MODELS[mid]}}


@app.api(name="status")
def api_status() -> dict:
    out, now = {}, time.time()
    for mid in MODELS:
        s = _st(mid)
        if mid in _pipes:
            s["state"] = "ready"
        busy = s["state"] in ("downloading", "loading")
        out[mid] = {"state": s["state"], "pct": s["pct"], "error": s["error"],
                    "elapsed": (now - s["t0"]) if busy and s["t0"] else s["took"],
                    "took": s["took"], "done": s["done"], "total": s["total"],
                    "speed": s["speed"], "eta": s["eta"]}
    return out


@app.api(name="load")
def api_load(model_id: str) -> dict:
    if model_id not in MODELS:
        return {"ok": False, "error": "نموذج غير معروف"}
    with _lock:
        if model_id not in _pipes and _st(model_id)["state"] not in ("downloading", "loading"):
            _st(model_id).update(state="downloading", t0=time.time())
            threading.Thread(target=_load_sync, args=(model_id,), daemon=True).start()
    return {"ok": True}


@app.api(name="unload")
def api_unload(model_id: str = "") -> dict:
    with _lock:
        for mid in ([model_id] if model_id else list(_pipes)):
            _pipes.pop(mid, None)
            if mid in _state:
                _state[mid].update(state="idle", pct=0)
        gc.collect()
    return {"ok": True}


@app.api(name="cancel")
def api_cancel(req_id: str) -> dict:
    _cancels[req_id] = True
    return {"ok": True}


@app.api(name="generate")
def api_generate(model_id: str, messages: str, req_id: str,
                 max_tokens: int = 256, temperature: float = 0.7, top_p: float = 0.9) -> dict:
    _cancels[req_id] = False
    try:
        p = _pipes.get(model_id)
        if p is None:
            yield {"ok": False, "error": "النموذج غير محمّل"}
            return
        kind, tok, model = MODELS[model_id]["kind"], p["tok"], p["model"]
        msgs = json.loads(messages)
        if len(msgs) > MAX_HISTORY + 1:  # نحتفظ بتعليمات النظام + آخر الرسائل
            head = msgs[:1] if msgs and msgs[0].get("role") == "system" else []
            msgs = head + msgs[-MAX_HISTORY:]

        special = True
        if kind == "chat":
            try:
                prompt = tok.apply_chat_template(msgs, tokenize=False,
                                                 add_generation_prompt=True, enable_thinking=False)
                special = False
            except Exception:
                prompt = _alpaca(msgs)
        elif kind == "t5":
            prompt = next((m["content"] for m in reversed(msgs) if m["role"] == "user"), "")
        else:
            prompt = _alpaca(msgs)

        enc = tok(prompt, return_tensors="pt", add_special_tokens=special)
        streamer = TextIteratorStreamer(tok, skip_prompt=(kind != "t5"), skip_special_tokens=True)
        kw = dict(**enc, max_new_tokens=int(max_tokens), streamer=streamer,
                  stopping_criteria=StoppingCriteriaList([StopOnCancel(req_id)]))
        if temperature > 0:
            kw.update(do_sample=True, temperature=float(temperature), top_p=float(top_p or 0.9))
        else:
            kw["do_sample"] = False
        if kind != "t5":
            kw.update(repetition_penalty=1.1, pad_token_id=tok.pad_token_id or tok.eos_token_id)

        errs = []

        def run():
            try:
                with torch.inference_mode():
                    model.generate(**kw)
            except Exception as e:
                errs.append(str(e))
                streamer.end()

        th = threading.Thread(target=run, daemon=True)
        t_start = time.time()
        th.start()
        out, pend, last = "", 0, t_start
        for piece in streamer:
            if _cancels.get(req_id):
                break
            out += piece
            pend += 1
            now = time.time()
            # تجميع التوكنات: نرسل فقط عند اكتمال دفعة أو مرور مهلة قصيرة
            if pend >= BATCH_TOKENS or now - last >= BATCH_SECONDS:
                yield {"ok": True, "text": out, "done": False}
                pend, last = 0, now
        th.join(timeout=2)
        if errs:
            yield {"ok": False, "error": errs[0]}
        else:
            secs = max(time.time() - t_start, 0.001)
            try:
                n_tok = len(tok.encode(out, add_special_tokens=False))
            except Exception:
                n_tok = 0
            yield {"ok": True, "text": out.strip(), "done": True,
                   "secs": round(secs, 2), "tokens": n_tok}
    except Exception as e:
        yield {"ok": False, "error": str(e)}
    finally:
        _cancels.pop(req_id, None)


@app.api(name="asr")
def api_asr(audio_b64: str) -> dict:
    try:
        p = _pipes.get("whisper")
        if p is None:
            return {"ok": False, "error": "نموذج الصوت غير محمّل"}
        data, sr = sf.read(io.BytesIO(base64.b64decode(audio_b64)), dtype="float32")
        if data.ndim > 1:
            data = data.mean(axis=1)
        if sr != 16000:
            data = np.interp(np.linspace(0, len(data) - 1, int(len(data) * 16000 / sr)),
                             np.arange(len(data)), data).astype("float32")
        feats = p["proc"](data[:16000 * 30], sampling_rate=16000, return_tensors="pt").input_features
        with torch.inference_mode():
            ids = p["model"].generate(feats, max_new_tokens=220)
        return {"ok": True, "text": p["proc"].batch_decode(ids, skip_special_tokens=True)[0].strip()}
    except Exception as e:
        return {"ok": False, "error": str(e)}


@app.api(name="caption")
def api_caption(image_b64: str) -> dict:
    try:
        p = _pipes.get("caption")
        if p is None:
            return {"ok": False, "error": "نموذج الصور غير محمّل"}
        img = Image.open(io.BytesIO(base64.b64decode(image_b64))).convert("RGB")
        pix = p["proc"](images=img, return_tensors="pt").pixel_values
        with torch.inference_mode():
            ids = p["model"].generate(pix, max_new_tokens=32, num_beams=4)
        return {"ok": True, "text": p["tok"].decode(ids[0], skip_special_tokens=True).strip()}
    except Exception as e:
        return {"ok": False, "error": str(e)}


@app.get("/", response_class=HTMLResponse)
async def serve_ui():
    with open(os.path.join(HERE, "index.html"), "r", encoding="utf-8") as f:
        return f.read()


app.launch(server_name="0.0.0.0",
           server_port=int(os.environ.get("GRADIO_SERVER_PORT", 7860)),
           ssr_mode=False, show_error=True)