Spaces:
Sleeping
Sleeping
| # -*- coding: utf-8 -*- | |
| """ | |
| 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 Ops: Builder (از golden_builder) + Cleaner/Deduper | |
| - Training: SFT/LoRA سبک روی causal + W&B logging/Artifacts | |
| - 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 | |
| import numpy as np | |
| import torch | |
| from torch.utils.data import Dataset | |
| from sklearn.model_selection import train_test_split | |
| import gradio as gr | |
| warnings.filterwarnings("ignore") | |
| # ====== ML & NLP ====== | |
| import transformers as tf | |
| from transformers import ( | |
| AutoTokenizer, AutoModelForCausalLM, | |
| Trainer, TrainingArguments, EarlyStoppingCallback | |
| ) | |
| # RAG stack | |
| import chromadb | |
| from rank_bm25 import BM25Okapi | |
| from sentence_transformers import CrossEncoder, SentenceTransformer, util as st_util | |
| # ========= Persian text 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) # ZWNJ | |
| s = re.sub(r"\s+", " ", s).strip() | |
| return s | |
| # ========================== | |
| # Configs | |
| # ========================== | |
| class ModelConfig: | |
| model_name: str = "Qwen/Qwen2.5-7B-Instruct" | |
| max_input_length: int = 4096 | |
| max_new_tokens: int = 512 | |
| temperature: float = 0.7 | |
| top_p: float = 0.9 | |
| do_sample: bool = True | |
| gradient_checkpointing: bool = True | |
| class RAGConfig: | |
| persist_dir: str = "./chroma_db" | |
| collection: str = "legal_articles" | |
| top_k: int = 8 | |
| similarity_threshold: float = 0.60 | |
| context_char_limit: int = 280 | |
| enable: bool = True | |
| reranker_name: str = "Alibaba-NLP/gte-multilingual-reranker-base" | |
| 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" # W&B | |
| max_grad_norm: float = 1.0 | |
| use_4bit: bool = True # QLoRA 4-bit (در صورت افزودن PEFT/TRL) | |
| max_seq_len: int = 2048 | |
| class SystemConfig: | |
| model: ModelConfig = field(default_factory=ModelConfig) | |
| rag: RAGConfig = field(default_factory=RAGConfig) | |
| train: TrainConfig = field(default_factory=TrainConfig) | |
| # ========================== | |
| # 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 bf16_supported(): | |
| return torch.cuda.is_available() and getattr(torch.cuda, "is_bf16_supported", lambda: False)() | |
| def log_deps(): | |
| try: | |
| 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) | |
| # ========================== | |
| # 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) | |
| 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) | |
| # reranker | |
| try: | |
| dev = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.reranker = CrossEncoder(self.cfg.reranker_name, device=dev) | |
| except Exception: | |
| self.reranker = None | |
| # BM25 | |
| 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() | |
| ids, docs, metas = [], [], [] | |
| 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 | |
| aid = str(obj.get(id_key, f"auto_{i}")) | |
| txt = normalize_fa(str(obj.get(text_key, "")).strip()) | |
| if not txt: continue | |
| ids.append(aid); docs.append(txt); metas.append({"article_id": aid}) | |
| if not ids: return "هیچ سندی برای ایندکس یافت نشد." | |
| self.collection.upsert(ids=ids, documents=docs, metadatas=metas) | |
| self._rebuild_bm25(ids, docs) | |
| return f"✅ {len(ids)} سند ایندکس شد (Dense+BM25)." | |
| def retrieve(self, query: str) -> List[Dict]: | |
| if not self.collection: return [] | |
| qn = normalize_fa(query) | |
| # Dense via Chroma | |
| 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}) | |
| # union by id | |
| 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] | |
| # threshold | |
| merged = [a for a in merged if a.get("similarity", 0) >= self.cfg.similarity_threshold] | |
| # rerank | |
| if self.reranker and merged: | |
| pairs = [(qn, a["text"]) for a in merged] | |
| 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) | |
| # ========================== | |
| 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 | |
| kwargs = {} | |
| 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 | |
| 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).to(self.model.device) | |
| with torch.no_grad(): | |
| out = self.model.generate( | |
| **enc, | |
| 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, | |
| ) | |
| 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() | |
| # ---------- W&B env ---------- | |
| 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, | |
| ) | |
| # Optional richer W&B init | |
| 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) | |
| # ========================== | |
| 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 | |
| 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): | |
| 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 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 [] | |
| 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 ans, refs | |
| 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)): | |
| 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 | |
| ) | |
| progress(0.95, desc="ذخیرهٔ آرتیفکتها") | |
| return f"✅ آموزش کامل شد و در {self.scfg.train.output_dir} ذخیره شد." | |
| # Dataset Builder (از ماژول شما) | |
| def build_dataset(self, raw_file, text_key: str, model_ckpt: str, batch_size: int, max_samples: int | 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}" | |
| # 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() | |
| try: | |
| print("[rag-bootstrap]", ensure_chroma_ready(self.scfg.rag.persist_dir, self.scfg.rag.collection), flush=True) | |
| except Exception as 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)") as app: | |
| 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 • Weight Tuning</p> | |
| </div> | |
| """) | |
| # --- Tab: Consultation --- | |
| 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 --- | |
| 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) | |
| # --- Tab: Dataset Builder --- | |
| 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) | |
| # --- Tab: Dataset Cleaning --- | |
| 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() | |
| # --- Tab: Training (W&B integrated) --- | |
| 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") | |
| # W&B controls | |
| 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) | |
| # --- 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] | |
| 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)) | |
| 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(lambda q, sys_p, rag, mnt, t, p: self.answer(q, sys_p, rag, mnt, t, p), | |
| inputs=[question, system_prompt, use_rag, max_new_tokens, temperature, top_p], | |
| outputs=[answer, refs]) | |
| index_btn.click(lambda f, ik, tk: self.build_index(f, ik, tk), | |
| inputs=[laws_file, id_key, text_key], outputs=index_status) | |
| build_btn.click(lambda rf, tk, ckpt, bs, mx: self.build_dataset(rf, tk, ckpt, bs, mx), | |
| inputs=[raw_file, ds_text_key, model_ckpt, ds_batch_size, max_samples], | |
| outputs=[out_file, build_status]) | |
| 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" | |
| train_btn.click( | |
| lambda prof, mid, files, rg, e, b, l, uw, wp, we, rn: | |
| self.train(_map_profile_to_id(prof, mid), files, rg, e, b, l, uw, wp, we, rn), | |
| 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 | |
| ) | |
| 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 | |
| # ========================== | |
| # Entrypoint | |
| # ========================== | |
| if __name__ == "__main__": | |
| app = LegalApp() | |
| ui = app.build_ui() | |
| try: | |
| ui = ui.queue() | |
| except TypeError: | |
| pass | |
| ui.launch(server_name="0.0.0.0", server_port=7860) | |