Update app.py

app.py

@@ -1,19 +1,21 @@
 # -*- coding: utf-8 -*-
 """
-Mahoon Legal AI — Causal-only Generation + Hybrid RAG + W&B-integrated Training
+Mahoon Legal AI — Causal-only Generation + Hybrid RAG + W&B Training + Weight Tuning
 - پاسخ‌زایی: Qwen2.5-7B, Llama-3.1-8B, Mistral-7B (همه causal)
 - RAG: Chroma + BM25 + CrossEncoder reranker (gte-multilingual-reranker-base)
-- Dataset: Builder (بر اساس golden_builder شما) + Cleaner/Deduper
+- Dataset Ops: Builder (از golden_builder) + Cleaner/Deduper
 - Training: SFT/LoRA سبک روی causal + W&B logging/Artifacts
-- UI: Gradio 5.47.0 (چهار تب: مشاوره، ایندکس قوانین، ساخت دیتاست، پاکسازی دیتاست، آموزش)
+- Tuning: Weight Tuning با W&B Sweep (weights_sweep.py)
+- UI: Gradio 5.47.0
 
+نکته: در Settings → Secrets مقدار `WANDB_API_KEY` را ست کنید (مقدار واقعی؛ placeholder 🟡 نگذارید).
 """
 
 from __future__ import annotations
 import os, sys, re, json, time, pickle, zipfile, warnings
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import List, Dict, Optional, Tuple
+from typing import List, Dict, Optional
 
 import numpy as np
 import torch
@@ -21,7 +23,6 @@ from torch.utils.data import Dataset
 from sklearn.model_selection import train_test_split
 
 import gradio as gr
-from packaging import version
 warnings.filterwarnings("ignore")
 
 # ====== ML & NLP ======
@@ -35,7 +36,6 @@ from transformers import (
 import chromadb
 from rank_bm25 import BM25Okapi
 from sentence_transformers import CrossEncoder, SentenceTransformer, util as st_util
-from langdetect import detect
 
 # ========= Persian text normalization =========
 ZWNJ = "\u200c"
@@ -50,7 +50,7 @@ def normalize_fa(s: str) -> str:
     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)    # نرمال‌سازی ZWNJ
+    s = re.sub(r"\s*‌\s*", ZWNJ, s)                      # ZWNJ
     s = re.sub(r"\s+", " ", s).strip()
     return s
 
@@ -79,9 +79,9 @@ class RAGConfig:
 
 @dataclass
 class TrainConfig:
-    base_model: str = "PartAI/Dorna-Llama3-8B-Instruct"   # قابل تغییر در UI
-    alt_model_1: str = "zpm/Llama-3.1-PersianQA"           # قابل تغییر در UI
-    hakim_model: str = "AI-Hoosh/HAKIM-7B"                 # به‌روزرسانی در UI
+    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
@@ -96,9 +96,9 @@ class TrainConfig:
     eval_strategy: str = "epoch"
     save_strategy: str = "epoch"
     save_total_limit: int = 2
-    report_to: str = "wandb"            # ← W&B
+    report_to: str = "wandb"            # W&B
     max_grad_norm: float = 1.0
-    use_4bit: bool = True               # QLoRA 4-bit
+    use_4bit: bool = True               # QLoRA 4-bit (در صورت افزودن PEFT/TRL)
     max_seq_len: int = 2048
 
 @dataclass
@@ -108,7 +108,7 @@ class SystemConfig:
     train: TrainConfig = field(default_factory=TrainConfig)
 
 # ==========================
-# Utils & deps logging
+# Helpers
 # ==========================
 def set_seed_all(seed: int = 42):
     import random
@@ -433,7 +433,6 @@ class TrainerManager:
             max_grad_norm=self.cfg.train.max_grad_norm,
         )
 
-        # ---------- Trainer + W&B callback ----------
         callbacks = [EarlyStoppingCallback(early_stopping_patience=2)]
         try:
             if use_wandb:
@@ -451,7 +450,7 @@ class TrainerManager:
             callbacks=callbacks,
         )
 
-        # Optional manual init for richer metadata
+        # Optional richer W&B init
         if use_wandb:
             try:
                 import wandb
@@ -474,7 +473,6 @@ class TrainerManager:
         trainer.save_model(self.cfg.train.output_dir)
         self.loader.tokenizer.save_pretrained(self.cfg.train.output_dir)
 
-        # Log artifacts to W&B
         if use_wandb:
             try:
                 import wandb
@@ -588,7 +586,7 @@ class LegalApp:
         set_seed_all(self.scfg.train.seed)
 
         progress(0.30, desc="آماده‌سازی دیتاست‌ها و RAG (اختیاری)")
-        out = tm.train_causal(
+        tm.train_causal(
             paths, use_rag=use_rag, use_wandb=use_wandb,
             wandb_project=wandb_project, wandb_entity=wandb_entity, run_name=run_name
         )
@@ -598,7 +596,10 @@ class LegalApp:
 
     # Dataset Builder (از ماژول شما)
     def build_dataset(self, raw_file, text_key: str, model_ckpt: str, batch_size: int, max_samples: int | None):
-        from golden_builder import load_json_or_jsonl, save_jsonl, GoldenBuilder
+        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:
@@ -613,6 +614,24 @@ class LegalApp:
         except Exception as e:
             return None, f"❌ خطا در ساخت دیتاست: {e}"
 
+    # Weight Tuning (W&B Sweep)
+    def run_weight_tune(self, f, tk, ms, runs, bs, proj, ent):
+        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}"
+
     # UI
     def build_ui(self):
         log_deps()
@@ -631,7 +650,7 @@ class LegalApp:
             gr.Markdown("""
             <div style='text-align:center;padding:18px'>
               <h1 style='margin-bottom:4px'>ماحون — Persian Legal (Causal-only)</h1>
-              <p style='color:#666'>Hybrid RAG • Qwen/Llama/Mistral • Dataset Ops • W&B Training</p>
+              <p style='color:#666'>Hybrid RAG • Qwen/Llama/Mistral • Dataset Ops • W&B Training • Weight Tuning</p>
             </div>
             """)
 
@@ -725,7 +744,7 @@ class LegalApp:
                 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("**راهنمای توکن W&B**: در Settings → Secrets مقدار `WANDB_API_KEY` را برابر با **🟡** قرار دهید.")
+                gr.Markdown("راهنما: در Settings → Secrets مقدار `WANDB_API_KEY` را تنظیم کنید (مقدار واقعی).")
 
                 train_files = gr.Files(label="JSONL Files", file_count="multiple", file_types=[".jsonl"])
                 with gr.Row():
@@ -735,6 +754,19 @@ class LegalApp:
                 train_btn = gr.Button("شروع آموزش", variant="primary")
                 train_status = gr.Textbox(label="وضعیت آموزش", interactive=False)
 
+            # --- Tab: Weight Tuning ---
+            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()
+
             # ---- Events ----
             def _resolve_gen(choice: str, override: str) -> str:
                 return override.strip() if override.strip() else default_gen_models[choice]
@@ -778,6 +810,26 @@ class LegalApp:
                 outputs=train_status
             )
 
+            clean_btn.click(
+                lambda f, th: (
+                    (lambda _p, _out:
+                        ( _out,
+                          f"✅ دیتاست پاک شد. تعداد رکوردهای نهایی: **{deduplicate_jsonl(_p, _out, sim_threshold=float(th))}**" )
+                    )(
+                        getattr(f, "name", None) or getattr(f, "path", None),
+                        f"/tmp/cleaned_{int(time.time())}.jsonl"
+                    ) if (getattr(f, 'name', None) or getattr(f, 'path', None)) else (None, "⚠️ فایل نامعتبر.")
+                ),
+                inputs=[raw_ds, sim_th],
+                outputs=[cleaned_out, clean_status]
+            )
+
+            run_tune.click(
+                lambda f, tk, ms, runs, bs, proj, ent: self.run_weight_tune(f, tk, ms, runs, bs, proj, ent),
+                inputs=[tune_file, tune_text_key, tune_max_samples, tune_runs, tune_batch, tune_proj, tune_entity],
+                outputs=tune_status
+            )
+
         return app
 
 # ==========================