Update app.py

app.py

@@ -1,208 +1,1033 @@
 # -*- coding: utf-8 -*-
-# Mahoon — Minimal RAG + Generation (ZeroGPU-ready, no training)
+"""
+Mahoon Legal AI — Causal-only Generation + Hybrid RAG + W&B + ZeroGPU + Role Gating
+
+- تب «مشاوره» برای همه تعاملی است.
+- تب‌های «ایندکس»، «ساخت دیتاست»، «پاکسازی»، «آموزش»، «Weight Tuning» برای بازدیدکننده فقط نمایشی‌اند؛
+  و سمت‌سرور نیز گِیت نقش دارد (ادمین/بازدیدکننده).
+
+پیش‌نیازها:
+- golden_builder.py , weights_sweep.py
+- Settings → Secrets: WANDB_API_KEY (در صورت استفاده از W&B)
+- Settings → Environment Variables: ADMIN_USERS (مثلاً: haji-mammad, teammate1)
+- requirements.txt (ZeroGPU-ready) شامل spaces>=0.42.0
+"""
+
+from __future__ import annotations
+
+# --- Telemetry hard-off + ZeroGPU SDK (must be before chroma import) ---
+import os, logging
+os.environ["CHROMA_TELEMETRY_ENABLED"] = "false"
+os.environ["ANONYMIZED_TELEMETRY"] = "false"
+
+import spaces  # ZeroGPU SDK
+
+# (اختیاری) کاهش نویز لاگ‌ها
+logging.getLogger("chromadb").setLevel(logging.ERROR)
+logging.getLogger("posthog").setLevel(logging.CRITICAL)
+# -----------------------------------------------------------------------
+
+import sys, re, json, time, pickle, zipfile, warnings
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Dict, Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+from sklearn.model_selection import train_test_split
 
-import os
-import json
 import gradio as gr
+warnings.filterwarnings("ignore")
 
-# =========================
-# ZeroGPU shim & marker
-# =========================
-try:
-    import spaces  # provided by HF Spaces runtime
-except Exception:
-    class _NoSpaces:
-        @staticmethod
-        def GPU(*a, **k):
-            def w(fn): return fn
-            return w
-    spaces = _NoSpaces()
-
-@spaces.GPU(duration=180)  # وجود این تابع جلوی ارور No @spaces.GPU را می‌گیرد
-def _zgpu_marker():
-    return "ok"
-
-# =========================
-# RAG (Chroma)
-# =========================
+# ====== Transformers ======
+import transformers as tf
+from transformers import (
+    AutoTokenizer, AutoModelForCausalLM,
+    Trainer, TrainingArguments, EarlyStoppingCallback
+)
+
+# ====== RAG stack ======
 import chromadb
 from chromadb.config import Settings
+from rank_bm25 import BM25Okapi
+from sentence_transformers import CrossEncoder, SentenceTransformer, util as st_util
+
+# ---- Monkeypatch Chroma telemetry (fallback) ----
+try:
+    import chromadb.telemetry as _ctel
+    try: _ctel.client = None
+    except Exception: pass
+    for _n in ("capture", "capture_event"):
+        if hasattr(_ctel, _n):
+            try: setattr(_ctel, _n, lambda *a, **k: None)
+            except Exception: pass
+    if hasattr(_ctel, "Telemetry"):
+        try: _ctel.Telemetry().capture = lambda *a, **k: None
+        except Exception: pass
+except Exception:
+    pass
+# -------------------------------------------------
+
+# ========= Persian normalization =========
+ZWNJ = "\u200c"
+AR_DIGITS = "٠١٢٣٤٥٦٧٨٩"
+FA_DIGITS = "۰۱۲۳۴۵۶۷۸۹"
+EN_DIGITS = "0123456789"
+
+def normalize_fa(s: str) -> str:
+    if not s:
+        return s
+    s = s.replace("\u064A", "ی").replace("\u0643", "ک")
+    s = re.sub(r"[\u064B-\u065F\u0610-\u061A]", "", s)
+    trans = {ord(a): e for a, e in zip(AR_DIGITS + FA_DIGITS, EN_DIGITS * 2)}
+    s = s.translate(trans)
+    s = re.sub(r"\s*‌\s*", ZWNJ, s)
+    s = re.sub(r"\s+", " ", s).strip()
+    return s
+
+# ==========================
+# Configs
+# ==========================
+@dataclass
+class ModelConfig:
+    model_name: str = "Qwen/Qwen2.5-7B-Instruct"
+    max_input_length: int = 3072
+    max_new_tokens: int = 256
+    temperature: float = 0.7
+    top_p: float = 0.9
+    do_sample: bool = True
+    gradient_checkpointing: bool = True
+
+@dataclass
+class RAGConfig:
+    persist_dir: str = "./chroma_db"
+    collection: str = "legal_articles"
+    top_k: int = 6
+    similarity_threshold: float = 0.68
+    context_char_limit: int = 260
+    enable: bool = True
+    reranker_name: str = "Alibaba-NLP/gte-multilingual-reranker-base"
+
+@dataclass
+class TrainConfig:
+    base_model: str = "PartAI/Dorna-Llama3-8B-Instruct"
+    alt_model_1: str = "zpm/Llama-3.1-PersianQA"
+    hakim_model: str = "AI-Hoosh/HAKIM-7B"
+    hooshvareh_model: str = "HooshvareLab/llama-fa-7b-instruct"
+    output_dir: str = "./mahoon_causal_lora"
+    seed: int = 42
+    test_size: float = 0.1
+    epochs: int = 2
+    batch_size: int = 2
+    grad_accum: int = 4
+    lr: float = 2e-4
+    warmup_ratio: float = 0.03
+    weight_decay: float = 0.0
+    logging_steps: int = 50
+    eval_strategy: str = "epoch"
+    save_strategy: str = "epoch"
+    save_total_limit: int = 2
+    report_to: str = "wandb"
+    max_grad_norm: float = 1.0
+    use_4bit: bool = False
+    max_seq_len: int = 2048
 
-CHROMA_DIR = os.environ.get("CHROMA_DIR", "./chroma_db")
-CHROMA_COLLECTION = os.environ.get("CHROMA_COLLECTION", "legal_articles")
+@dataclass
+class SystemConfig:
+    model: ModelConfig = field(default_factory=ModelConfig)
+    rag: RAGConfig = field(default_factory=RAGConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
 
-def _norm_id(x: str) -> str:
-    x = (x or "").replace("\u064A","ی").replace("\u0643","ک")
-    trans = {ord(a): b for a,b in zip("٠١٢٣٤٥٦٧٨٩۰۱۲۳۴۵۶۷۸۹","01234567890123456789")}
-    return "".join((x.translate(trans))).replace(" ", "")
+# ==========================
+# Helpers
+# ==========================
+def set_seed_all(seed: int = 42):
+    import random
+    random.seed(seed); np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 
-def build_rag():
-    client = chromadb.PersistentClient(
-        path=CHROMA_DIR,
-        settings=Settings(anonymized_telemetry=False)
-    )
+def bf16_supported():
+    return torch.cuda.is_available() and getattr(torch.cuda, "is_bf16_supported", lambda: False)()
+
+def log_deps():
     try:
-        col = client.get_or_create_collection(CHROMA_COLLECTION)
-    except Exception:
-        col = client.get_collection(CHROMA_COLLECTION)
-    return col
+        import accelerate, datasets
+        print("[deps]",
+              f"python={sys.version.split()[0]}",
+              f"transformers={tf.__version__}",
+              f"accelerate={accelerate.__version__}",
+              f"datasets={datasets.__version__}",
+              f"gradio={gr.__version__}",
+              flush=True)
+    except Exception as e:
+        print("[deps] warn:", e, flush=True)
 
-def retrieve(col, query: str, top_k: int, thr: float):
+# ==========================
+# Role gating helpers
+# ==========================
+def _get_username(request: gr.Request) -> str | None:
     try:
-        res = col.query(
-            query_texts=[query],
-            n_results=int(top_k),
-            include=["documents","metadatas","distances"]
-        )
-        docs = res.get("documents",[[]])[0]
-        metas= res.get("metadatas",[[]])[0]
-        dists= res.get("distances",[[]])[0]
-        out=[]
-        for i,(d,m,dist) in enumerate(zip(docs, metas, dists)):
-            sim = 1.0 - float(dist)
-            if sim >= float(thr):
-                out.append({
-                    "article_id": _norm_id((m or {}).get("article_id", f"unk_{i}")),
-                    "text": d,
-                    "similarity": sim
-                })
-        return out
+        return getattr(request, "username", None)
     except Exception:
-        return []
-
-def build_context(arts, limit=320):
-    if not arts: return ""
-    bullets = [f"• ماده {a['article_id']}: {a['text'][:limit]}..." for a in arts]
-    return "مواد مرتبط:\n" + "\n".join(bullets)
-
-# =========================
-# Generation (Transformers)
-# =========================
-# برای اجتناب از نیاز زودهنگام به torch، import را داخل توابع انجام می‌دهیم.
-MODEL_CHOICES = {
-    "Qwen2.5-7B Instruct":       "Qwen/Qwen2.5-7B-Instruct",
-    "Llama 3.2 3B Instruct":     "meta-llama/Llama-3.2-3B-Instruct",
-    "Mistral 7B Instruct v0.2":  "mistralai/Mistral-7B-Instruct-v0.2"
-}
-DEFAULT_MODEL_KEY = os.environ.get("DEFAULT_MODEL_KEY", "Llama 3.2 3B Instruct")
-
-_loader = {"tk": None, "model_id": None}
-_rag_col = None
-
-def lazy_bootstrap(selected_key: str):
-    """اتصال به ایندکس RAG + Warm tokenizer. وزن مدل را بعداً در تابع GPU لود می‌کنیم."""
-    global _rag_col, _loader
-    # RAG
-    if _rag_col is None:
-        try:
-            _rag_col = build_rag()
-        except Exception as e:
-            return f"❌ خطا در اتصال RAG: {e}"
-
-    # Tokenizer
-    wanted = MODEL_CHOICES.get(selected_key, MODEL_CHOICES[DEFAULT_MODEL_KEY])
-    if _loader["model_id"] != wanted or _loader["tk"] is None:
-        from transformers import AutoTokenizer
-        tk = AutoTokenizer.from_pretrained(wanted)
-        if tk.pad_token is None and tk.eos_token:
-            tk.pad_token = tk.eos_token
-        _loader.update({"tk": tk, "model_id": wanted})
-
-    return f"✅ آماده · ایندکس: {CHROMA_COLLECTION} · مدل: {wanted}"
-
-def _format_prompt(context: str, question: str) -> str:
-    if context:
-        return f"{context}\nسوال: {question}\nپاسخ:"
-    return f"سوال: {question}\nپاسخ:"
-
-@spaces.GPU(duration=240)
-def answer_gpu(model_key, question, use_rag, top_k, thr, max_new_tokens, temperature, top_p):
-    """اینفرنس روی GPU (ZeroGPU per-call)."""
-    try:
-        if not question or not question.strip():
-            return "لطفاً سؤال را وارد کنید.", ""
+        return None
+
+def is_admin(request: gr.Request) -> bool:
+    uname = _get_username(request)
+    if not uname:
+        return False
+    author = os.getenv("SPACE_AUTHOR_NAME", "").strip()
+    allow = {u.strip() for u in os.getenv("ADMIN_USERS", "").split(",") if u.strip()}
+    return (uname == author) or (uname in allow)
+
+# ==========================
+# RAG: Chroma + BM25 + CrossEncoder reranker
+# ==========================
+class LegalRAG:
+    def __init__(self, cfg: RAGConfig):
+        self.cfg = cfg
+        self.client = None
+        self.collection = None
+        self.reranker: Optional[CrossEncoder] = None
+        self.bm25 = None
+        self.bm25_ids: List[str] = []
+        self.bm25_path = str(Path(self.cfg.persist_dir) / "bm25.pkl")
+
+    def init(self):
+        Path(self.cfg.persist_dir).mkdir(parents=True, exist_ok=True)
+        self.client = chromadb.PersistentClient(
+            path=self.cfg.persist_dir,
+            settings=Settings(anonymized_telemetry=False)
+        )
+        try:
+            self.collection = self.client.get_or_create_collection(self.cfg.collection)
+        except Exception:
+            try: self.collection = self.client.get_collection(self.cfg.collection)
+            except Exception: self.collection = self.client.create_collection(self.cfg.collection)
+
+        try:
+            self.reranker = CrossEncoder(self.cfg.reranker_name, device="cpu")
+        except Exception:
+            self.reranker = None
+
+        if Path(self.bm25_path).exists():
+            with open(self.bm25_path, "rb") as f:
+                obj = pickle.load(f)
+            self.bm25 = obj["bm25"]; self.bm25_ids = obj["ids"]
+
+    def _rebuild_bm25(self, ids: List[str], docs: List[str]):
+        corpus = [normalize_fa(d).split() for d in docs]
+        self.bm25 = BM25Okapi(corpus)
+        self.bm25_ids = ids
+        with open(self.bm25_path, "wb") as f:
+            pickle.dump({"bm25": self.bm25, "ids": self.bm25_ids}, f)
+
+    def index_jsonl(self, jsonl_path: str, id_key="article_id", text_key="text"):
+        if not self.collection: self.init()
+
+        seen: Dict[str, int] = {}
+        ids, docs, metas = [], [], []
+
+        def _norm_id(x: str) -> str:
+            x = x or ""
+            x = x.replace("\u064A", "ی").replace("\u0643", "ک")
+            trans = {ord(a): e for a, e in zip("٠١٢٣٤٥٦٧٨٩۰۱۲۳۴۵۶۷۸۹", "01234567890123456789")}
+            x = x.translate(trans)
+            x = re.sub(r"\s+", "", x)
+            return x
+
+        with open(jsonl_path, "r", encoding="utf-8") as f:
+            for i, line in enumerate(f):
+                s = line.strip()
+                if not s: continue
+                try: obj = json.loads(s)
+                except: continue
+
+                raw_id = str(obj.get(id_key, f"auto_{i}"))
+                base_id = _norm_id(raw_id)
+                txt = normalize_fa(str(obj.get(text_key, "")).strip())
+                if not txt: continue
+
+                if base_id in seen:
+                    seen[base_id] += 1
+                    uid = f"{base_id}__d{seen[base_id]}"
+                    dupe_idx = seen[base_id]
+                else:
+                    seen[base_id] = 1
+                    uid = base_id
+                    dupe_idx = 1
+
+                ids.append(uid); docs.append(txt); metas.append({"article_id": base_id, "dupe_idx": dupe_idx})
+
+        if not ids:
+            return "هیچ سندی برای ایندکس یافت نشد."
+
+        self.collection.upsert(ids=ids, documents=docs, metadatas=metas)
+        self._rebuild_bm25(ids, docs)
+
+        dup_count = sum(1 for _, c in seen.items() if c > 1)
+        return f"✅ {len(ids)} سند ایندکس شد (Dense+BM25). شناسه‌های تکراری: {dup_count} کلید (با پسوند __dN یکتا شدند)."
+
+    def retrieve(self, query: str) -> List[Dict]:
+        if not self.collection: return []
+        qn = normalize_fa(query)
+
+        # Dense
+        try:
+            res = self.collection.query(
+                query_texts=[qn],
+                n_results=max(self.cfg.top_k * 3, 20),
+                include=["documents", "metadatas", "distances"],
+            )
+            out = []
+            docs = res.get("documents", [[]])[0]
+            metas = res.get("metadatas", [[]])[0]
+            dists = res.get("distances", [[1.0]])[0]
+            for i, (doc, meta, dist) in enumerate(zip(docs, metas, dists)):
+                sim = 1.0 - float(dist)
+                out.append({"article_id": (meta or {}).get("article_id", f"unk_{i}"),
+                            "text": doc, "similarity": sim})
+        except Exception:
+            out = []
+
+        # BM25
+        bm25_hits = []
+        if self.bm25 is not None and self.bm25_ids:
+            scores = self.bm25.get_scores(normalize_fa(qn).split())
+            idxs = np.argsort(scores)[::-1][:max(self.cfg.top_k * 3, 20)]
+            smax = float(scores.max() + 1e-8)
+            for j in idxs:
+                aid = self.bm25_ids[int(j)]
+                try:
+                    got = self.collection.get(ids=[aid])
+                    tdoc = got["documents"][0]
+                except Exception:
+                    tdoc = ""
+                bm25_hits.append({"article_id": aid, "text": tdoc, "similarity": float(scores[j]) / smax})
+
+        # merge
+        pool: Dict[str, Dict] = {}
+        for a in out + bm25_hits:
+            if a["article_id"] not in pool or a.get("similarity", 0) > pool[a["article_id"]].get("similarity", 0):
+                pool[a["article_id"]] = a
+        merged = [a for a in pool.values() if a.get("text") and len(a["text"]) > 15]
+        merged = [a for a in merged if a.get("similarity", 0) >= self.cfg.similarity_threshold]
+
+        # rerank (GPU only during predict)
+        if merged and self.reranker:
+            pairs = [(qn, a["text"]) for a in merged]
+            try:
+                with spaces.GPU(duration=30):
+                    scores = self.reranker.predict(pairs)
+            except Exception:
+                scores = self.reranker.predict(pairs)
+            for a, s in zip(merged, scores): a["score"] = float(s)
+            merged = sorted(merged, key=lambda x: x.get("score", 0), reverse=True)[: self.cfg.top_k]
+        else:
+            merged = sorted(merged, key=lambda x: x.get("similarity", 0), reverse=True)[: self.cfg.top_k]
+        return merged
+
+    def build_context(self, arts: List[Dict]) -> str:
+        if not arts: return ""
+        bullets = [f"• ماده {a['article_id']}: {a['text'][:self.cfg.context_char_limit]}..." for a in arts]
+        return "مواد مرتبط:\n" + "\n".join(bullets)
+
+# ========= RAG bootstrap from repo =========
+def parse_law_textfile_to_jsonl(txt_path: str, out_jsonl: str):
+    pat = re.compile(r"(?:ماده|مادّه)\s+(\d+)\s*[:\-–]\s*(.+)")
+    rows = []
+    with open(txt_path, "r", encoding="utf-8") as f:
+        for line in f:
+            s = line.strip()
+            if not s: continue
+            m = pat.match(s)
+            if not m: continue
+            aid = m.group(1); body = m.group(2).strip()
+            if len(body) < 12: continue
+            rows.append({"article_id": aid, "text": normalize_fa(body)})
+    if not rows: raise RuntimeError("هیچ ماده‌ای با الگوی تعریف‌شده پیدا نشد.")
+    with open(out_jsonl, "w", encoding="utf-8") as g:
+        for r in rows: g.write(json.dumps(r, ensure_ascii=False) + "\n")
+    return len(rows)
+
+def ensure_chroma_ready(persist_dir="./chroma_db", collection="legal_articles") -> str:
+    Path(persist_dir).mkdir(parents=True, exist_ok=True)
+    if any(Path(persist_dir).glob("*")):
+        return f"ChromaDB موجود است."
+    zip_path = Path("./chroma_legal_db.zip")
+    if zip_path.exists():
+        try:
+            with zipfile.ZipFile(zip_path, "r") as z: z.extractall(persist_dir)
+            return "ChromaDB از zip بازیابی شد."
+        except Exception: pass
+    txt_path = Path("./all_legal_sentences.txt")
+    if txt_path.exists():
+        n = parse_law_textfile_to_jsonl(str(txt_path), "./laws.jsonl")
+        rag_local = LegalRAG(RAGConfig(persist_dir=persist_dir, collection=collection))
+        rag_local.init()
+        msg = rag_local.index_jsonl("./laws.jsonl", id_key="article_id", text_key="text")
+        return f"از متن خام {n} رکورد استخراج شد. {msg}"
+    return "پایگاه RAG موجود نیست و منبع خامی هم برای ساخت پیدا نشد."
+
+# ==========================
+# Loader + Generator (Causal-only, ZeroGPU)
+# ==========================
+class CausalLoader:
+    def __init__(self, mcfg: ModelConfig):
+        self.cfg = mcfg
+        self.tokenizer = None
+        self.model = None
+
+    def load(self, model_name: str):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+        if self.tokenizer.pad_token is None and hasattr(self.tokenizer, "eos_token"):
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+
+        try:
+            with spaces.GPU(duration=90):
+                kwargs = {"low_cpu_mem_usage": True}
+                if torch.cuda.is_available():
+                    kwargs["device_map"] = "auto"
+                    kwargs["torch_dtype"] = torch.bfloat16 if bf16_supported() else torch.float16
+                self.model = AutoModelForCausalLM.from_pretrained(model_name, **kwargs)
+                if self.cfg.gradient_checkpointing and hasattr(self.model, "gradient_checkpointing_enable"):
+                    try: self.model.gradient_checkpointing_enable()
+                    except Exception: pass
+        except Exception:
+            self.model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
+
+        return self
+
+class Generator:
+    def __init__(self, loader: CausalLoader, mcfg: ModelConfig):
+        self.tk = loader.tokenizer
+        self.model = loader.model
+        self.cfg = mcfg
+
+    def generate(self, question: str, context: str = "", system_prompt: str = "You are a helpful Persian legal assistant.") -> str:
+        parts = []
+        if system_prompt: parts.append(f"<|system|>\n{system_prompt}")
+        if context: parts.append(f"<|system|>\nاز منابع زیر استفاده کن و استنادی پاسخ بده:\n{context}")
+        parts.append(f"<|user|>\n{question}")
+        prompt = "\n".join(parts) + "\n<|assistant|>\n"
+
+        enc = self.tk(prompt, return_tensors="pt", truncation=True, max_length=self.cfg.max_input_length)
+
+        try:
+            with spaces.GPU(duration=60):
+                dev_model = next(self.model.parameters()).device if hasattr(self.model, "parameters") else "cpu"
+                inputs = {k: v.to(dev_model) for k, v in enc.items()}
+                with torch.no_grad():
+                    out = self.model.generate(
+                        **inputs,
+                        max_new_tokens=self.cfg.max_new_tokens,
+                        do_sample=self.cfg.do_sample,
+                        temperature=self.cfg.temperature,
+                        top_p=self.cfg.top_p,
+                        pad_token_id=self.tk.pad_token_id or self.tk.eos_token_id,
+                    )
+        except Exception:
+            inputs = {k: v for k, v in enc.items()}
+            with torch.no_grad():
+                out = self.model.generate(
+                    **inputs,
+                    max_new_tokens=min(self.cfg.max_new_tokens, 256),
+                    do_sample=self.cfg.do_sample,
+                    temperature=self.cfg.temperature,
+                    top_p=self.cfg.top_p,
+                    pad_token_id=self.tk.pad_token_id or self.tk.eos_token_id,
+                )
+
+        return self.tk.decode(out[0], skip_special_tokens=True)
+
+# ==========================
+# Datasets & Trainer (Causal-only, W&B)
+# ==========================
+def read_jsonl_files(paths: List[str]) -> List[Dict]:
+    data: List[Dict] = []
+    for p in paths:
+        if not p: continue
+        with open(p, 'r', encoding='utf-8') as f:
+            for line in f:
+                s = line.strip()
+                if not s: continue
+                try: data.append(json.loads(s))
+                except json.JSONDecodeError: continue
+    return data
+
+class CausalJSONLDataset(Dataset):
+    def __init__(self, data: List[Dict], tokenizer, max_len: int, rag: Optional[LegalRAG] = None, enhance_every:int = 8):
+        self.tk = tokenizer
+        self.max_len = max_len
+        self.items = []
+        for i, ex in enumerate(data):
+            src = normalize_fa(str(ex.get("input", "")).strip())
+            tgt = normalize_fa(str(ex.get("output", "")).strip())
+            if not src or not tgt: continue
+            ctx = ""
+            if rag and i % enhance_every == 0:
+                arts = rag.retrieve(src)
+                ctx = rag.build_context(arts)
+            text = ""
+            if ctx: text += f"<|system|>\nاز منابع زیر استفاده کن:\n{ctx}\n"
+            text += f"<|system|>\nYou are a helpful Persian legal assistant.\n"
+            text += f"<|user|>\n{src}\n<|assistant|>\n{tgt}"
+            self.items.append(text)
+
+    def __len__(self): return len(self.items)
+
+    def __getitem__(self, idx):
+        text = self.items[idx]
+        enc = self.tk(text, max_length=self.max_len, padding="max_length", truncation=True)
+        input_ids = torch.tensor(enc["input_ids"])
+        attn = torch.tensor(enc["attention_mask"])
+        labels = input_ids.clone(); labels[attn == 0] = -100
+        return {"input_ids": input_ids, "attention_mask": attn, "labels": labels}
+
+def safe_training_args(**kwargs):
+    return TrainingArguments(**kwargs)
+
+class TrainerManager:
+    def __init__(self, syscfg: SystemConfig, loader: CausalLoader):
+        self.cfg = syscfg
+        self.loader = loader
+
+    def train_causal(self, train_paths: List[str], use_rag: bool = True, use_wandb: bool = True,
+                     wandb_project: str = "mahoon-legal-ai", wandb_entity: str = "", run_name: str = "mahoon_causal_lora"):
+        set_seed_all(self.cfg.train.seed)
+        data = read_jsonl_files(train_paths)
+        train, val = train_test_split(data, test_size=self.cfg.train.test_size, random_state=self.cfg.train.seed)
+
+        rag = LegalRAG(self.cfg.rag) if (use_rag and self.cfg.rag.enable) else None
+        if rag: rag.init()
+
+        ds_tr = CausalJSONLDataset(train, self.loader.tokenizer, self.cfg.train.max_seq_len, rag)
+        ds_va = CausalJSONLDataset(val,   self.loader.tokenizer, self.cfg.train.max_seq_len, None)
+
+        fp16_ok = torch.cuda.is_available() and not bf16_supported()
+        bf16_ok = bf16_supported()
+
+        if use_wandb:
+            os.environ.setdefault("WANDB_PROJECT", wandb_project or "mahoon-legal-ai")
+            if wandb_entity: os.environ.setdefault("WANDB_ENTITY", wandb_entity)
+            os.environ.pop("WANDB_DISABLED", None)
+        else:
+            os.environ["WANDB_DISABLED"] = "true"
 
+        args = safe_training_args(
+            output_dir=self.cfg.train.output_dir,
+            num_train_epochs=self.cfg.train.epochs,
+            learning_rate=self.cfg.train.lr,
+            per_device_train_batch_size=self.cfg.train.batch_size,
+            per_device_eval_batch_size=self.cfg.train.batch_size,
+            gradient_accumulation_steps=self.cfg.train.grad_accum,
+            warmup_ratio=self.cfg.train.warmup_ratio,
+            weight_decay=self.cfg.train.weight_decay,
+            evaluation_strategy=self.cfg.train.eval_strategy,
+            save_strategy=self.cfg.train.save_strategy,
+            save_total_limit=self.cfg.train.save_total_limit,
+            load_best_model_at_end=True,
+            metric_for_best_model="eval_loss",
+            logging_steps=self.cfg.train.logging_steps,
+            report_to=(["wandb"] if use_wandb else ["none"]),
+            run_name=run_name,
+            fp16=fp16_ok, bf16=bf16_ok,
+            max_grad_norm=self.cfg.train.max_grad_norm,
+        )
+
+        callbacks = [EarlyStoppingCallback(early_stopping_patience=2)]
+        try:
+            if use_wandb:
+                from transformers.integrations import WandbCallback
+                callbacks.append(WandbCallback())
+        except Exception:
+            pass
+
+        trainer = Trainer(
+            model=self.loader.model,
+            args=args,
+            train_dataset=ds_tr,
+            eval_dataset=ds_va,
+            tokenizer=self.loader.tokenizer,
+            callbacks=callbacks,
+        )
+
+        if use_wandb:
+            try:
+                import wandb
+                wandb.init(project=os.getenv("WANDB_PROJECT", "mahoon-legal-ai"),
+                           entity=os.getenv("WANDB_ENTITY"),
+                           name=run_name,
+                           config={
+                               "base_model": self.loader.model.name_or_path,
+                               "epochs": self.cfg.train.epochs,
+                               "batch": self.cfg.train.batch_size,
+                               "grad_accum": self.cfg.train.grad_accum,
+                               "lr": self.cfg.train.lr,
+                               "max_seq_len": self.cfg.train.max_seq_len,
+                               "use_rag": use_rag,
+                           })
+            except Exception:
+                pass
+
+        trainer.train()
+        trainer.save_model(self.cfg.train.output_dir)
+        self.loader.tokenizer.save_pretrained(self.cfg.train.output_dir)
+
+        if use_wandb:
+            try:
+                import wandb
+                art = wandb.Artifact("mahoon-model", type="model")
+                art.add_dir(self.cfg.train.output_dir)
+                wandb.log_artifact(art)
+                wandb.finish()
+            except Exception:
+                pass
+
+# ==========================
+# Dataset utilities (Cleaner/Deduper)
+# ==========================
+def deduplicate_jsonl(in_path: str, out_path: str, sim_threshold: float = 0.90, text_keys=("input","output")) -> int:
+    rows = []
+    with open(in_path, "r", encoding="utf-8") as f:
+        for line in f:
+            s = line.strip()
+            if not s: continue
+            try: obj = json.loads(s)
+            except: continue
+            for k in text_keys:
+                if k in obj: obj[k] = normalize_fa(str(obj[k]))
+            rows.append(obj)
+    if not rows: raise RuntimeError("هیچ رکورد معتبری در ورودی نبود.")
+    model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+    embs = model.encode([r.get("input","") for r in rows], convert_to_tensor=True, show_progress_bar=False, normalize_embeddings=True)
+    kept, seen = [], torch.zeros(len(rows), dtype=torch.bool)
+    for i in range(len(rows)):
+        if seen[i]: continue
+        sims = st_util.cos_sim(embs[i], embs)[0]
+        dup_idx = (sims >= sim_threshold).nonzero(as_tuple=True)[0].tolist()
+        for j in dup_idx: seen[j] = True
+        kept.append(rows[i])
+    with open(out_path, "w", encoding="utf-8") as g:
+        for r in kept: g.write(json.dumps(r, ensure_ascii=False) + "\n")
+    return len(kept)
+
+# ==========================
+# App (Gradio) + Role Gating
+# ==========================
+class LegalApp:
+    def __init__(self, scfg: Optional[SystemConfig] = None):
+        self.scfg = scfg or SystemConfig()
+        self.rag = LegalRAG(self.scfg.rag)
+        self.loader: Optional[CausalLoader] = None
+        self.gen: Optional[Generator] = None
+
+    def _file_paths(self, files: List[gr.File]) -> List[str]:
+        paths = []
+        for f in (files or []):
+            p = getattr(f, "name", None) or getattr(f, "path", None)
+            if p: paths.append(p)
+        return paths
+
+    # Core (مشاوره/لود آزاد است)
+    def load(self, model_name: str):
+        self.loader = CausalLoader(self.scfg.model).load(model_name)
+        self.gen = Generator(self.loader, self.scfg.model)
         # RAG
-        arts = retrieve(_rag_col, question, int(top_k), float(thr)) if use_rag else []
-        ctx  = build_context(arts) if arts else ""
-        prompt = _format_prompt(ctx, question)
-
-        # بارگذاری وزن‌ها روی GPU رزروشده
-        from transformers import AutoTokenizer, AutoModelForCausalLM
-        model_id = _loader["model_id"] or MODEL_CHOICES.get(model_key) or MODEL_CHOICES[DEFAULT_MODEL_KEY]
-        tk = _loader["tk"] or AutoTokenizer.from_pretrained(model_id)
-        mdl = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")  # ZeroGPU → GPU attach
-
-        enc = tk(prompt, return_tensors="pt")
-        enc = {k: v.to(mdl.device) for k,v in enc.items()}
-        out = mdl.generate(
-            **enc,
-            max_new_tokens=int(max_new_tokens),
-            do_sample=True,
-            temperature=float(temperature),
-            top_p=float(top_p),
-            pad_token_id=tk.pad_token_id or tk.eos_token_id
+        msg_rag = "RAG غیرفعال"
+        if self.scfg.rag.enable:
+            try:
+                self.rag = LegalRAG(self.scfg.rag); self.rag.init()
+                msg_rag = "RAG آماده است"
+            except Exception as e:
+                msg_rag = f"RAG خطا: {e}"
+        return f"مدل بارگذاری شد: {model_name}\n{msg_rag}"
+
+    # --- گیت سمت‌سرور: فقط ادمین ---
+    def build_index(self, laws_file: gr.File, id_key: str, text_key: str, request: gr.Request):
+        if not is_admin(request):
+            return "🔒 این عملیات فقط برای مدیران فعال است."
+        if not self.scfg.rag.enable: return "RAG غیرفعال است."
+        try:
+            self.rag.init()
+            p = getattr(laws_file, "name", None) or getattr(laws_file, "path", None)
+            if not p: return "فایل قوانین معتبر نیست."
+            return self.rag.index_jsonl(p, id_key=id_key, text_key=text_key)
+        except Exception as e:
+            return f"خطا در ایندکس: {e}"
+
+    def build_dataset(self, raw_file, text_key: str, model_ckpt: str, batch_size: int, max_samples: int | None, request: gr.Request):
+        if not is_admin(request):
+            return None, "🔒 این عملیات فقط برای مدیران فعال است."
+        try:
+            from golden_builder import load_json_or_jsonl, save_jsonl, GoldenBuilder
+        except Exception as e:
+            return None, f"❌ golden_builder.py یافت نشد/قابل import نیست: {e}"
+        path = getattr(raw_file, "name", None) or getattr(raw_file, "path", None)
+        if not path: return None, "⚠️ فایل ورودی معتبر نیست."
+        try:
+            data = load_json_or_jsonl(path)
+            if max_samples and int(max_samples) > 0: data = data[:int(max_samples)]
+            gb = GoldenBuilder(model_name=model_ckpt)
+            rows = gb.build(data, text_key=text_key, batch_size=int(batch_size))
+            out_dir = "/tmp/mahoon_datasets"; Path(out_dir).mkdir(parents=True, exist_ok=True)
+            out_path = f"{out_dir}/golden_{os.path.basename(path)}.jsonl"
+            save_jsonl(rows, out_path)
+            return out_path, f"✅ {len(rows)} رکورد تولید شد."
+        except Exception as e:
+            return None, f"❌ خطا در ساخت دیتاست: {e}"
+
+    def train(self, model_name: str, files: List[gr.File], use_rag: bool, epochs: int, batch: int, lr: float,
+              use_wandb: bool, wandb_project: str, wandb_entity: str, run_name: str,
+              progress=gr.Progress(track_tqdm=True), request: gr.Request = None):
+        if not is_admin(request):
+            return "🔒 این عملیات فقط برای مدیران فعال است."
+        progress(0.05, desc="راه‌اندازی")
+        self.scfg.train.epochs = int(epochs)
+        self.scfg.train.batch_size = int(batch)
+        self.scfg.train.lr = float(lr)
+
+        progress(0.10, desc="بارگذاری مدل/توکنایزر")
+        self.loader = CausalLoader(self.scfg.model).load(model_name)
+
+        paths = self._file_paths(files)
+        if not paths: return "⚠️ هیچ فایل JSONL برای آموزش انتخاب نشده."
+
+        tm = TrainerManager(self.scfg, self.loader)
+        set_seed_all(self.scfg.train.seed)
+
+        progress(0.30, desc="آماده‌سازی دیتاست‌ها و RAG (اختیاری)")
+        tm.train_causal(
+            paths, use_rag=use_rag, use_wandb=use_wandb,
+            wandb_project=wandb_project, wandb_entity=wandb_entity, run_name=run_name
         )
-        text = tk.decode(out[0], skip_special_tokens=True)
+
+        progress(0.95, desc="ذخیرهٔ آرتیفکت‌ها")
+        return f"✅ آموزش کامل شد و در {self.scfg.train.output_dir} ذخیره شد."
+
+    def run_weight_tune(self, f, tk, ms, runs, bs, proj, ent, request: gr.Request):
+        if not is_admin(request):
+            return "🔒 این عملیات فقط برای مدیران فعال است."
+        p = getattr(f, "name", None) or getattr(f, "path", None)
+        if not p:
+            return "⚠️ فایل داده نامعتبر است."
+        try:
+            from weights_sweep import run_sweep
+        except Exception as e:
+            return f"❌ weights_sweep.py یافت نشد/قابل import نیست: {e}"
+        os.environ.setdefault("WANDB_PROJECT", proj or "mahoon-legal-ai")
+        if ent: os.environ.setdefault("WANDB_ENTITY", ent)
+        try:
+            run_sweep(data_path=p, text_key=tk, max_samples=int(ms), batch_size=int(bs),
+                      project=proj, entity=ent, count=int(runs))
+            return "✅ Sweep اجرا شد. بهترین Run را در W&B بررسی و وزن‌ها را تثبیت کنید."
+        except Exception as e:
+            return f"❌ خطا در اجرای Sweep: {e}"
+
+    def apply_best_weights(self, wandb_project: str, wandb_entity: str, metric: str = "pass_rate", request: gr.Request = None):
+        if request is not None and not is_admin(request):
+            return "🔒 این عملیات فقط برای مدیران فعال است."
+        try:
+            import wandb, json as _json
+        except Exception as e:
+            return f"❌ W&B در محیط در دسترس نیست: {e}"
+
+        try:
+            api = wandb.Api()
+            proj_path = f"{wandb_entity}/{wandb_project}" if wandb_entity else wandb_project
+            runs = api.runs(proj_path, filters={"state": "finished"})
+        except Exception as e:
+            return f"❌ عدم دسترسی به پروژه W&B ({wandb_project}): {e}"
+
+        best_run = None; best_val = float("-inf")
+        for r in runs:
+            s = r.summary or {}
+            if "weights" in s and metric in s:
+                try: val = float(s[metric])
+                except Exception: continue
+                if val > best_val: best_val, best_run = val, r
+
+        if not best_run:
+            return "⚠️ هیچ Run واجد شرایطی با summary['weights'] و متریک موردنظر پیدا نشد."
+
+        weights = best_run.summary.get("weights", {})
+        if not isinstance(weights, dict) or not weights:
+            return "⚠️ فرمت وزن‌های بهترین Run نامعتبر است."
+
+        try:
+            with open("legal_entity_weights.json", "w", encoding="utf-8") as f:
+                _json.dump(weights, f, ensure_ascii=False, indent=2)
+        except Exception as e:
+            return f"❌ خطا در نوشتن legal_entity_weights.json: {e}"
+
+        rid = getattr(best_run, "id", "unknown")
+        return f"✅ وزن‌ها اعمال شد از Run `{rid}` با {metric}={best_val:.4f}. فایل: `legal_entity_weights.json`"
+
+    # Consultation (عمومی)
+    def answer(self, question: str, system_prompt: str, use_rag: bool, max_new_tokens: int, temperature: float, top_p: float):
+        if not question.strip(): return "لطفاً سوال خود را وارد کنید.", ""
+        if not self.gen: return "ابتدا مدل را بارگذاری کنید.", ""
+        self.scfg.model.max_new_tokens = int(max_new_tokens)
+        self.scfg.model.temperature = float(temperature)
+        self.scfg.model.top_p = float(top_p)
+
+        arts = self.rag.retrieve(question) if (use_rag and self.scfg.rag.enable and self.rag.collection) else []
+        max_refs = 4
+        if arts: arts = arts[:max_refs]
+        ctx = self.rag.build_context(arts) if arts else ""
+        ans = self.gen.generate(question, ctx, system_prompt)
 
         refs = ""
         if arts:
-            refs = "\n\n" + "\n".join([
-                f"**ماده {a['article_id']}** (شباهت: {a['similarity']:.2f})\n{a['text'][:380]}..."
-                for a in arts
-            ])
-        return text, refs
-    except Exception as e:
-        return f"❌ خطای اینفرنس: {e}", ""
-
-# =========================
-# UI (Gradio 5.47)
-# =========================
-with gr.Blocks(title="Mahoon — Minimal RAG+Gen", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""
-    <div style='text-align:center;padding:14px'>
-      <h2 style='margin:0'>ماحون (مینیمال) — پاسخ حقوقی با RAG</h2>
-      <p style='color:#666'>اینفرنس ZeroGPU · ایندکس آماده · بدون آموزش</p>
-    </div>
-    """)
-
-    with gr.Row():
-        model_dd = gr.Dropdown(choices=list(MODEL_CHOICES.keys()),
-                               value=DEFAULT_MODEL_KEY,
-                               label="مدل تولید")
-        use_rag  = gr.Checkbox(value=True, label="استفاده از RAG؟")
-        top_k    = gr.Slider(1, 10, value=5, step=1, label="Top-K")
-        thr      = gr.Slider(0.50, 0.95, value=0.60, step=0.01, label="آستانه شباهت")
-
-    with gr.Accordion("پارامترهای تولید", open=False):
-        max_new_tokens = gr.Slider(64, 1024, value=256, step=16, label="max_new_tokens")
-        temperature    = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="temperature")
-        top_p          = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
-
-    question = gr.Textbox(lines=3, label="سؤال")
-    ask_btn  = gr.Button("پرسش", variant="primary")
-    answer   = gr.Markdown(label="پاسخ")
-    refs     = gr.Markdown(label="مواد مرتبط")
-
-    status = gr.Markdown("⏳ آماده‌سازی…")
-
-    def _warmup(mkey):
-        try:
-            return lazy_bootstrap(mkey)
+            refs = "\n\n" + "\n".join([f"**ماده {a['article_id']}** (شباهت: {a.get('similarity',0):.2f})\n{a['text'][:320]}..." for a in arts])
+        return ans, refs
+
+    # UI
+    def build_ui(self):
+        log_deps()
+        try:
+            print("[rag-bootstrap]", ensure_chroma_ready(self.scfg.rag.persist_dir, self.scfg.rag.collection), flush=True)
         except Exception as e:
-            return f"❌ Bootstrap error: {e}"
+            print("[rag-bootstrap] error:", e, flush=True)
+
+        default_gen_models = {
+            "Qwen2.5-7B Instruct": "Qwen/Qwen2.5-7B-Instruct",
+            "Llama-3.1-8B Instruct": "meta-llama/Llama-3.1-8B-Instruct",
+            "Mistral-7B Instruct (v0.3)": "mistralai/Mistral-7B-Instruct-v0.3",
+        }
+
+        with gr.Blocks(title="ماحون — مشاور حقوقی (Causal-only, ZeroGPU)") as app:
+            # بنر نقش
+            role_banner = gr.Markdown()
+
+            gr.Markdown("""
+            <div style='text-align:center;padding:18px'>
+              <h1 style='margin-bottom:4px'>ماحون — Persian Legal (Causal-only, ZeroGPU)</h1>
+              <p style='color:#666'>Hybrid RAG • Qwen/Llama/Mistral • Dataset Ops • W&B Training • Weight Tuning</p>
+            </div>
+            """)
+
+            # --- Tab: Consultation (interactive for all) ---
+            with gr.Tab("مشاوره"):
+                with gr.Row():
+                    gen_model_dd = gr.Dropdown(choices=list(default_gen_models.keys()), value="Qwen2.5-7B Instruct", label="مدل تولید")
+                    gen_model_id = gr.Textbox(value=default_gen_models["Qwen2.5-7B Instruct"], label="Model ID (قابل ویرایش)")
+                with gr.Row():
+                    use_rag = gr.Checkbox(value=True, label="RAG فعال باشد؟")
+                    persist_dir = gr.Textbox(value=self.scfg.rag.persist_dir, label="مسیر ChromaDB")
+                    collection = gr.Textbox(value=self.scfg.rag.collection, label="نام کالکشن")
+                with gr.Row():
+                    top_k = gr.Slider(1, 15, value=self.scfg.rag.top_k, step=1, label="Top-K")
+                    threshold = gr.Slider(0.3, 0.95, value=self.scfg.rag.similarity_threshold, step=0.01, label="آستانه شباهت")
+                load_btn = gr.Button("بارگذاری مدل", variant="primary")
+                status = gr.Textbox(label="وضعیت", interactive=False)
+
+                with gr.Accordion("پارامترهای تولید", open=False):
+                    system_prompt = gr.Textbox(value="You are a helpful Persian legal assistant.", label="System prompt")
+                    max_new_tokens = gr.Slider(64, 2048, value=self.scfg.model.max_new_tokens, step=16, label="max_new_tokens")
+                    temperature = gr.Slider(0.0, 1.5, value=self.scfg.model.temperature, step=0.05, label="temperature")
+                    top_p = gr.Slider(0.1, 1.0, value=self.scfg.model.top_p, step=0.05, label="top_p")
+
+                question = gr.Textbox(lines=3, label="سوال حقوقی")
+                gr.Examples(
+                    examples=[
+                        ["در صورت نقض قرارداد EPC چه راهکارهای حقوقی دارم؟"],
+                        ["آیا درج شرط عدم رقابت در قرارداد کار قانونی است؟"],
+                        ["حق و حقوق کارگر در صورت اخراج فوری چیست؟"],
+                    ],
+                    inputs=question, label="نمونه پرسش‌ها"
+                )
+                ask_btn = gr.Button("پرسش", variant="primary")
+                answer = gr.Markdown(label="پاسخ"); refs = gr.Markdown(label="مواد قانونی مرتبط")
+
+            # --- Tab: Indexing (view-only for visitors) ---
+            with gr.Tab("ایندکس قوانین"):
+                gr.Markdown("فایل JSONL قوانین را بارگذاری و ایندکس کنید (کلیدها: `article_id`, `text`).")
+                laws_file = gr.File(label="فایل JSONL قوانین", file_types=[".jsonl"])
+                id_key = gr.Textbox(value="article_id", label="کلید شناسه ماده")
+                text_key = gr.Textbox(value="text", label="کلید متن ماده")
+                index_btn = gr.Button("ایندکس‌سازی قوانین"); index_status = gr.Textbox(label="وضعیت ایندکس", interactive=False)
+                index_widgets = [laws_file, id_key, text_key, index_btn]
+
+            # --- Tab: Dataset Builder (view-only for visitors) ---
+            with gr.Tab("ساخت دیتاست"):
+                gr.Markdown("فایل خام (JSON/JSONL) → خروجی JSONL سازگار با `{input, output}` (از golden_builder).")
+                raw_file = gr.File(label="فایل خام", file_types=[".json",".jsonl"])
+                with gr.Row():
+                    ds_text_key = gr.Textbox(value="متن_کامل", label="کلید متن (text_key)")
+                    model_ckpt = gr.Dropdown(
+                        choices=["google/mt5-base", "google/flan-t5-base", "t5-base"],
+                        value="google/mt5-base",
+                        label="مدل خلاصه‌ساز برای ساخت دیتاست (فقط Builder)"
+                    )
+                with gr.Row():
+                    ds_batch_size = gr.Slider(1, 16, value=4, step=1, label="Batch size")
+                    max_samples = gr.Number(value=0, label="حداکثر نمونه (۰=همه)")
+                build_btn = gr.Button("ساخت دیتاست", variant="primary")
+                out_file = gr.File(label="دانلود خروجی JSONL", interactive=False)
+                build_status = gr.Textbox(label="وضعیت", interactive=False)
+                builder_widgets = [raw_file, ds_text_key, model_ckpt, ds_batch_size, max_samples, build_btn]
+
+            # --- Tab: Dataset Cleaning (view-only for visitors) ---
+            with gr.Tab("پاکسازی دیتاست"):
+                gr.Markdown("نرمال‌سازی فارسی + حذف تکراری‌های معنایی (cosine). ورودی: JSONL `{input, output}`.")
+                raw_ds = gr.File(label="JSONL ورودی", file_types=[".jsonl"])
+                sim_th = gr.Slider(0.80, 0.98, value=0.90, step=0.01, label="آستانه شباهت (cosine)")
+                clean_btn = gr.Button("اجرای پاکسازی", variant="primary")
+                cleaned_out = gr.File(label="دانلود JSONL پاک", interactive=False)
+                clean_status = gr.Markdown()
+                clean_widgets = [raw_ds, sim_th, clean_btn]
+
+            # --- Tab: Training (view-only for visitors) ---
+            with gr.Tab("آموزش"):
+                gr.Markdown("SFT/LoRA روی مدل‌های causal (فقط `{input, output}`) + W&B logging.")
+                with gr.Row():
+                    model_train_dd = gr.Dropdown(
+                        choices=[
+                            "HAKIM (Editable ID below)",
+                            "Hooshvareh (Editable ID below)",
+                            "Dorna-Llama3-8B",
+                            "PersianQA-8B",
+                            "Custom (Editable ID below)"
+                        ],
+                        value="HAKIM (Editable ID below)", label="پروفایل مدل"
+                    )
+                    model_train_id = gr.Textbox(value="AI-Hoosh/HAKIM-7B", label="HF Model ID (قابل ویرایش)")
+                use_rag_train = gr.Checkbox(value=True, label="RAG-enhanced Training")
+
+                use_wandb = gr.Checkbox(value=True, label="W&B logging فعال باشد؟")
+                wandb_project = gr.Textbox(value="mahoon-legal-ai", label="WANDB_PROJECT")
+                wandb_entity = gr.Textbox(value="", label="WANDB_ENTITY (اختیاری)")
+                run_name = gr.Textbox(value="mahoon_causal_lora", label="Run name")
+                gr.Markdown("راهنما: در Settings → Secrets مقدار `WANDB_API_KEY` را تنظیم کنید (مقدار واقعی).")
+
+                train_files = gr.Files(label="JSONL Files", file_count="multiple", file_types=[".jsonl"])
+                with gr.Row():
+                    epochs = gr.Slider(1, 6, value=2, step=1, label="epochs")
+                    batch = gr.Slider(1, 8, value=2, step=1, label="batch per device")
+                    lr = gr.Number(value=2e-4, label="learning rate")
+                train_btn = gr.Button("شروع آموزش", variant="primary")
+                train_status = gr.Textbox(label="وضعیت آموزش", interactive=False)
+                train_widgets = [model_train_dd, model_train_id, use_rag_train, use_wandb, wandb_project, wandb_entity,
+                                 run_name, train_files, epochs, batch, lr, train_btn]
+
+            # --- Tab: Weight Tuning (view-only for visitors) ---
+            with gr.Tab("Weight Tuning"):
+                gr.Markdown("تیون خودکار وزن‌های موجودیت با W&B Sweep. ابتدا در Settings→Secrets مقدار `WANDB_API_KEY` را ست کنید.")
+                tune_file = gr.File(label="فایل داده (JSON/JSONL)", file_types=[".json",".jsonl"])
+                tune_text_key = gr.Textbox(value="متن_کامل", label="کلید متن")
+                tune_max_samples = gr.Slider(50, 400, value=120, step=10, label="حداکثر نمونه")
+                tune_runs = gr.Slider(4, 64, value=16, step=4, label="تعداد ران Sweep")
+                tune_batch = gr.Slider(1, 4, value=2, step=1, label="batch size Builder")
+                tune_proj = gr.Textbox(value="mahoon-legal-ai", label="WANDB_PROJECT")
+                tune_entity = gr.Textbox(value="", label="WANDB_ENTITY (اختیاری)")
+                run_tune = gr.Button("شروع Sweep", variant="primary")
+                tune_status = gr.Markdown()
+
+                gr.Markdown("---")
+                gr.Markdown("اعمال خودکار بهترین وزن‌ها از داشبورد W&B (بر اساس بالاترین `pass_rate`).")
+                metric_dd = gr.Dropdown(choices=["pass_rate"], value="pass_rate", label="متریک انتخاب بهترین Run")
+                apply_btn = gr.Button("اعمال بهترین وزن‌ها از W&B", variant="secondary")
+                tuning_widgets = [tune_file, tune_text_key, tune_max_samples, tune_runs, tune_batch,
+                                  tune_proj, tune_entity, run_tune, metric_dd, apply_btn]
+
+            # ---- Events (مشاوره آزاد / عملیاتِ ادمینی با گیت) ----
+            def _resolve_gen(choice: str, override: str) -> str:
+                return override.strip() if override.strip() else default_gen_models[choice]
+
+            def _on_load(choice, override, rag, pdir, coll, k, th):
+                self.scfg.rag.enable = bool(rag)
+                self.scfg.rag.persist_dir = pdir
+                self.scfg.rag.collection = coll
+                self.scfg.rag.top_k = int(k)
+                self.scfg.rag.similarity_threshold = float(th)
+                return self.load(_resolve_gen(choice, override))
+
+            def _whoami(request: gr.Request):
+                u = _get_username(request) or "Visitor"
+                return f"👤 کاربر: **{u}** — دسترسی: {'مدیریتی' if is_admin(request) else 'بازدیدکننده (فقط مشاهده)'}"
+
+            load_btn.click(_on_load,
+                           inputs=[gen_model_dd, gen_model_id, use_rag, persist_dir, collection, top_k, threshold],
+                           outputs=status)
+
+            ask_btn.click(self.answer,
+                          inputs=[question, system_prompt, use_rag, max_new_tokens, temperature, top_p],
+                          outputs=[answer, refs])
+
+            # ادمینی: استفاده از request injection (Gradio به‌طور خودکار تزریق می‌کند)
+            def _index_handler(f, ik, tk, request: gr.Request):
+                return self.build_index(f, ik, tk, request)
+            index_btn.click(_index_handler, inputs=[laws_file, id_key, text_key], outputs=index_status)
+
+            def _build_ds_handler(rf, tk, ckpt, bs, mx, request: gr.Request):
+                return self.build_dataset(rf, tk, ckpt, bs, mx, request)
+            build_btn.click(_build_ds_handler,
+                            inputs=[raw_file, ds_text_key, model_ckpt, ds_batch_size, max_samples],
+                            outputs=[out_file, build_status])
+
+            def _train_handler(prof, mid, files, rg, e, b, l, uw, wp, we, rn, request: gr.Request):
+                def _map_profile_to_id(profile: str, current_id: str) -> str:
+                    if current_id.strip(): return current_id.strip()
+                    if "Dorna" in profile: return "PartAI/Dorna-Llama3-8B-Instruct"
+                    if "PersianQA" in profile: return "zpm/Llama-3.1-PersianQA"
+                    if "HAKIM" in profile: return "AI-Hoosh/HAKIM-7B"
+                    if "Hooshvareh" in profile: return "HooshvareLab/llama-fa-7b-instruct"
+                    return "PartAI/Dorna-Llama3-8B-Instruct"
+                model_id = _map_profile_to_id(prof, mid)
+                return self.train(model_id, files, rg, e, b, l, uw, wp, we, rn, request=request)
+            train_btn.click(_train_handler,
+                            inputs=[model_train_dd, model_train_id, train_files, use_rag_train, epochs, batch, lr,
+                                    use_wandb, wandb_project, wandb_entity, run_name],
+                            outputs=train_status)
+
+            def _clean_handler(f, th):
+                p = getattr(f, "name", None) or getattr(f, "path", None)
+                if not p: return None, "⚠️ فایل نامعتبر."
+                outp = f"/tmp/cleaned_{int(time.time())}.jsonl"
+                n = deduplicate_jsonl(p, outp, sim_threshold=float(th))
+                return outp, f"✅ دیتاست پاک شد. تعداد رکوردهای نهایی: **{n}**"
+            clean_btn.click(_clean_handler, inputs=[raw_ds, sim_th], outputs=[cleaned_out, clean_status])
+
+            def _tune_handler(f, tk, ms, runs, bs, proj, ent, request: gr.Request):
+                return self.run_weight_tune(f, tk, ms, runs, bs, proj, ent, request)
+            run_tune.click(_tune_handler,
+                           inputs=[tune_file, tune_text_key, tune_max_samples, tune_runs, tune_batch, tune_proj, tune_entity],
+                           outputs=tune_status)
+
+            def _apply_best_handler(proj, ent, m, request: gr.Request):
+                return self.apply_best_weights(proj, ent, m, request)
+            apply_btn.click(_apply_best_handler,
+                            inputs=[tune_proj, tune_entity, metric_dd],
+                            outputs=tune_status)
+
+            # --- Lock non-consultation tabs for visitors on load ---
+            def _gate_all(request: gr.Request):
+                admin = is_admin(request)
+                role_txt = f"👤 کاربر: **{_get_username(request) or 'Visitor'}** — دسترسی: {'مدیریتی' if admin else 'بازدیدکننده (فقط مشاهده)'}"
+                if not admin:
+                    lock = gr.update(interactive=False)
+                    updates = [lock] * (len(index_widgets) + len(builder_widgets) + len(clean_widgets) + len(train_widgets) + len(tuning_widgets))
+                else:
+                    unlock = gr.update(interactive=True)
+                    updates = [unlock] * (len(index_widgets) + len(builder_widgets) + len(clean_widgets) + len(train_widgets) + len(tuning_widgets))
+                return [role_txt] + updates
 
-    demo.load(_warmup, inputs=[model_dd], outputs=status)
+            app.load(_whoami, inputs=None, outputs=role_banner)
+            app.load(_gate_all, inputs=None,
+                     outputs=[role_banner] + index_widgets + builder_widgets + clean_widgets + train_widgets + tuning_widgets)
 
-    ask_btn.click(
-        answer_gpu,
-        inputs=[model_dd, question, use_rag, top_k, thr, max_new_tokens, temperature, top_p],
-        outputs=[answer, refs]
-    )
+        return app
 
+# ==========================
+# Entrypoint
+# ==========================
 if __name__ == "__main__":
+    app = LegalApp()
+    ui = app.build_ui()
     try:
-        demo = demo.queue()  # پایدارتر روی Gradio 5.x
+        ui = ui.queue()  # پایدار برای ZeroGPU
     except TypeError:
         pass
-    demo.launch(ssr_mode=False)
+    ui.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)