Update app.py

app.py

@@ -1,20 +1,208 @@
-# app.py (smoke test)
-import gradio as gr, os
+# -*- coding: utf-8 -*-
+# Mahoon — Minimal RAG + Generation (ZeroGPU-ready, no training)
 
-def ping(q: str):
-    return f"pong · echo: {q or ''}"
+import os
+import json
+import gradio as gr
 
-with gr.Blocks(title="Smoke Test") as demo:
-    gr.Markdown("✅ UI up. If you see this, Gradio is fine.")
-    inp = gr.Textbox(label="input")
-    out = gr.Textbox(label="output")
-    btn = gr.Button("Ping")
-    btn.click(ping, inp, out)
+# =========================
+# 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)
+# =========================
+import chromadb
+from chromadb.config import Settings
+
+CHROMA_DIR = os.environ.get("CHROMA_DIR", "./chroma_db")
+CHROMA_COLLECTION = os.environ.get("CHROMA_COLLECTION", "legal_articles")
+
+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(" ", "")
+
+def build_rag():
+    client = chromadb.PersistentClient(
+        path=CHROMA_DIR,
+        settings=Settings(anonymized_telemetry=False)
+    )
+    try:
+        col = client.get_or_create_collection(CHROMA_COLLECTION)
+    except Exception:
+        col = client.get_collection(CHROMA_COLLECTION)
+    return col
+
+def retrieve(col, query: str, top_k: int, thr: float):
+    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
+    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 "لطفاً سؤال را وارد کنید.", ""
+
+        # 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
+        )
+        text = tk.decode(out[0], skip_special_tokens=True)
+
+        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)
+        except Exception as e:
+            return f"❌ Bootstrap error: {e}"
+
+    demo.load(_warmup, inputs=[model_dd], outputs=status)
+
+    ask_btn.click(
+        answer_gpu,
+        inputs=[model_dd, question, use_rag, top_k, thr, max_new_tokens, temperature, top_p],
+        outputs=[answer, refs]
+    )
 
 if __name__ == "__main__":
-    # روی Spaces نیازی به پارامتر خاصی نیست؛ SSR را هم خاموش نگه‌دار که safe باشد
     try:
-        demo = demo.queue()  # پایدارتر روی 5.x
+        demo = demo.queue()  # پایدارتر روی Gradio 5.x
     except TypeError:
         pass
     demo.launch(ssr_mode=False)