app(7).py

# -*- coding: utf-8 -*-
"""
Mahoun — Ultimate Legal AI (Single-File, Modular, Polished UI)
هستهٔ جدید ماحون با ادغام اجزای قبلی (RAG پیشرفته + Training برای Seq2Seq و Causal) و UI زیباتر.

ویژگی‌ها:
- Multi-Architecture: "seq2seq" (T5/MT5/FLAN-T5) و "causal" (Mistral/LLaMA).
- Loader/Generator یکپارچه + Prompt تطبیقی برحسب معماری.
- RAG پیشرفته با ChromaDB (پیکربندی مسیر، نام کالکشن، top_k، threshold، قطع متن).
- Training کامل برای هر دو معماری (Trainer, EarlyStopping, bf16/fp16, gradient_accumulation).
- Gradio UI بازطراحی‌شده (تم تمیز، کارت‌ها، مثال‌ها، وضعیت زنده، کنترل‌های تولید، انتخاب مدل/معماری/دیتابیس).

حداقل نیازمندی‌ها (requirements.txt):
transformers>=4.44.0
sentencepiece
accelerate
bitsandbytes
chromadb
sentence-transformers
scikit-learn
gradio
torch>=2.1
"""
from __future__ import annotations
import os, json, gc, warnings, textwrap
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Dict, Optional, Tuple

import torch
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split

from transformers import (
    AutoTokenizer,
    AutoModelForSeq2SeqLM,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    EarlyStoppingCallback,
    DataCollatorForSeq2Seq,
)

import chromadb
from sentence_transformers import SentenceTransformer
import gradio as gr

warnings.filterwarnings("ignore")

# ==========================
# Config
# ==========================
@dataclass
class ModelConfig:
    model_name: str = "google/mt5-base"
    architecture: str = "seq2seq"  # "seq2seq" | "causal"
    max_input_length: int = 1024
    max_target_length: int = 512
    max_new_tokens: int = 384
    temperature: float = 0.7
    top_p: float = 0.9
    num_beams: int = 4

@dataclass
class RAGConfig:
    embedding_model: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
    persist_dir: str = "./chroma_db"
    collection: str = "legal_articles"
    top_k: int = 5
    similarity_threshold: float = 0.66   # 0..1 (بزرگ‌تر=سخت‌گیرتر)
    context_char_limit: int = 300        # حداکثر کاراکتر هر ماده در Context

@dataclass
class TrainConfig:
    output_dir: str = "./mahoon_model"
    seed: int = 42
    test_size: float = 0.1
    epochs: int = 2
    batch_size: int = 2
    grad_accum: int = 2
    lr: float = 3e-5
    use_bf16: bool = True

@dataclass
class SystemConfig:
    model: ModelConfig = field(default_factory=ModelConfig)
    rag: RAGConfig = field(default_factory=RAGConfig)
    train: TrainConfig = field(default_factory=TrainConfig)

# ==========================
# Utils
# ==========================
def set_seed_all(seed: int = 42):
    import random
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


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:
                    obj = json.loads(s)
                    data.append(obj)
                except json.JSONDecodeError:
                    continue
    return data

# ==========================
# RAG
# ==========================
class LegalRAG:
    def __init__(self, cfg: RAGConfig):
        self.cfg = cfg
        self.client = None
        self.collection = None
        self.embedder: Optional[SentenceTransformer] = None

    def init(self):
        Path(self.cfg.persist_dir).mkdir(parents=True, exist_ok=True)
        self.client = chromadb.PersistentClient(path=self.cfg.persist_dir)
        # get_or_create برای سازگاری نسخه‌های مختلف chroma
        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)
        self.embedder = SentenceTransformer(self.cfg.embedding_model)

    def retrieve(self, query: str) -> List[Dict]:
        if not self.collection:
            return []
        try:
            res = self.collection.query(
                query_texts=[query],
                n_results=self.cfg.top_k,
                include=["documents","metadatas","distances"],
            )
            out = []
            for i,(doc, meta, dist) in enumerate(zip(res.get('documents',[['']])[0], res.get('metadatas',[['']])[0], res.get('distances',[[1.0]])[0])):
                sim = 1 - float(dist)
                if sim >= self.cfg.similarity_threshold:
                    out.append({
                        "article_id": (meta or {}).get("article_id", f"unk_{i}"),
                        "text": doc,
                        "similarity": sim,
                    })
            return out
        except Exception:
            return []

    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)

# ==========================
# Loader + Generator
# ==========================
class ModelLoader:
    def __init__(self, mcfg: ModelConfig):
        self.cfg = mcfg
        self.tokenizer = None
        self.model = None

    def load(self):
        self.tokenizer = AutoTokenizer.from_pretrained(self.cfg.model_name)
        dtype = torch.bfloat16 if torch.cuda.is_available() else None
        if self.cfg.architecture == "seq2seq":
            self.model = AutoModelForSeq2SeqLM.from_pretrained(
                self.cfg.model_name, device_map="auto" if torch.cuda.is_available() else None, torch_dtype=dtype
            )
        elif self.cfg.architecture == "causal":
            self.model = AutoModelForCausalLM.from_pretrained(
                self.cfg.model_name, device_map="auto" if torch.cuda.is_available() else None, torch_dtype=dtype
            )
            if self.tokenizer.pad_token is None and hasattr(self.tokenizer, 'eos_token'):
                self.tokenizer.pad_token = self.tokenizer.eos_token
        else:
            raise ValueError("Unsupported architecture")
        return self

class Generator:
    def __init__(self, loader: ModelLoader, mcfg: ModelConfig):
        self.tk = loader.tokenizer
        self.model = loader.model
        self.cfg = mcfg

    def generate(self, question: str, context: str = "") -> str:
        if self.cfg.architecture == "seq2seq":
            inp = f"{context}\nسوال: {question}" if context else f"سوال: {question}"
            enc = self.tk(inp, return_tensors="pt", truncation=True, max_length=self.cfg.max_input_length)
            enc = {k: v.to(self.model.device) for k,v in enc.items()}
            out = self.model.generate(
                **enc,
                max_length=self.cfg.max_target_length,
                num_beams=self.cfg.num_beams,
                early_stopping=True,
            )
        else:  # causal
            prompt = f"{context}\nسوال: {question}\nپاسخ:" if context else f"سوال: {question}\nپاسخ:"
            enc = self.tk(prompt, return_tensors="pt", truncation=True, max_length=self.cfg.max_input_length)
            enc = {k: v.to(self.model.device) for k,v in enc.items()}
            out = self.model.generate(
                **enc,
                max_new_tokens=self.cfg.max_new_tokens,
                do_sample=True,
                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
# ==========================
class Seq2SeqJSONLDataset(Dataset):
    def __init__(self, data: List[Dict], tokenizer, max_inp: int, max_tgt: int, rag: Optional[LegalRAG] = None, enhance_every:int = 10):
        self.tk = tokenizer
        self.max_inp = max_inp
        self.max_tgt = max_tgt
        self.items = []
        for i, ex in enumerate(data):
            src = str(ex.get("input", "")).strip()
            tgt = str(ex.get("output", "")).strip()
            if not src or not tgt:
                continue
            inp = src
            if rag and i % enhance_every == 0:
                arts = rag.retrieve(src)
                ctx = rag.build_context(arts)
                if ctx:
                    inp = f"<CONTEXT>{ctx}</CONTEXT>\n<QUESTION>{src}</QUESTION>"
            self.items.append((inp, tgt))

    def __len__(self):
        return len(self.items)

    def __getitem__(self, idx):
        inp, tgt = self.items[idx]
        model_inputs = self.tk(inp, max_length=self.max_inp, padding="max_length", truncation=True)
        labels = self.tk(text_target=tgt, max_length=self.max_tgt, padding="max_length", truncation=True)
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

class CausalJSONLDataset(Dataset):
    def __init__(self, data: List[Dict], tokenizer, max_inp: int, rag: Optional[LegalRAG] = None, enhance_every:int = 10):
        self.tk = tokenizer
        self.max_inp = max_inp
        self.items = []
        for i, ex in enumerate(data):
            src = str(ex.get("input", "")).strip()
            tgt = 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 = f"{ctx}\nسوال: {src}\nپاسخ: {tgt}" if ctx else f"سوال: {src}\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_inp, padding="max_length", truncation=True)
        input_ids = torch.tensor(enc["input_ids"])
        return {"input_ids": input_ids, "attention_mask": torch.tensor(enc["attention_mask"]), "labels": input_ids.clone()}

# ==========================
# Trainer Manager
# ==========================
class TrainerManager:
    def __init__(self, syscfg: SystemConfig, loader: ModelLoader):
        self.cfg = syscfg
        self.loader = loader

    def train_seq2seq(self, train_paths: List[str], use_rag: bool = True):
        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 else None
        if rag:
            rag.init()
        ds_tr = Seq2SeqJSONLDataset(train, self.loader.tokenizer, self.cfg.model.max_input_length, self.cfg.model.max_target_length, rag)
        ds_va = Seq2SeqJSONLDataset(val,   self.loader.tokenizer, self.cfg.model.max_input_length, self.cfg.model.max_target_length, None)
        collator = DataCollatorForSeq2Seq(tokenizer=self.loader.tokenizer, model=self.loader.model)
        fp16_ok = torch.cuda.is_available() and (not self.cfg.train.use_bf16)
        bf16_ok = torch.cuda.is_available() and self.cfg.train.use_bf16
        args = TrainingArguments(
            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=0.05,
            weight_decay=0.01,
            evaluation_strategy="epoch",
            save_strategy="epoch",
            save_total_limit=2,
            load_best_model_at_end=True,
            metric_for_best_model="eval_loss",
            predict_with_generate=True,
            generation_max_length=self.cfg.model.max_target_length,
            generation_num_beams=self.cfg.model.num_beams,
            logging_steps=50,
            report_to="none",
            fp16=fp16_ok,
            bf16=bf16_ok,
        )
        trainer = Trainer(
            model=self.loader.model,
            args=args,
            train_dataset=ds_tr,
            eval_dataset=ds_va,
            data_collator=collator,
            tokenizer=self.loader.tokenizer,
            callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
        )
        trainer.train()
        trainer.save_model(self.cfg.train.output_dir)
        self.loader.tokenizer.save_pretrained(self.cfg.train.output_dir)

    def train_causal(self, train_paths: List[str], use_rag: bool = True):
        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 else None
        if rag:
            rag.init()
        ds_tr = CausalJSONLDataset(train, self.loader.tokenizer, self.cfg.model.max_input_length, rag)
        ds_va = CausalJSONLDataset(val,   self.loader.tokenizer, self.cfg.model.max_input_length, None)
        fp16_ok = torch.cuda.is_available() and (not self.cfg.train.use_bf16)
        bf16_ok = torch.cuda.is_available() and self.cfg.train.use_bf16
        args = TrainingArguments(
            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=0.05,
            weight_decay=0.01,
            evaluation_strategy="epoch",
            save_strategy="epoch",
            save_total_limit=2,
            load_best_model_at_end=True,
            metric_for_best_model="eval_loss",
            logging_steps=50,
            report_to="none",
            fp16=fp16_ok,
            bf16=bf16_ok,
        )
        trainer = Trainer(
            model=self.loader.model,
            args=args,
            train_dataset=ds_tr,
            eval_dataset=ds_va,
            tokenizer=self.loader.tokenizer,
            callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
        )
        trainer.train()
        trainer.save_model(self.cfg.train.output_dir)
        self.loader.tokenizer.save_pretrained(self.cfg.train.output_dir)

# ==========================
# 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[ModelLoader] = None
        self.gen: Optional[Generator] = None

    # --- core actions ---
    def load(self, model_name: str, arch: str, use_rag: bool, persist_dir: str, collection: str, top_k: int, threshold: float):
        # configure
        self.scfg.model.model_name = model_name
        self.scfg.model.architecture = arch
        self.scfg.rag.persist_dir = persist_dir
        self.scfg.rag.collection = collection
        self.scfg.rag.top_k = int(top_k)
        self.scfg.rag.similarity_threshold = float(threshold)
        # load model
        self.loader = ModelLoader(self.scfg.model).load()
        self.gen = Generator(self.loader, self.scfg.model)
        # load rag
        msg_rag = "RAG غیر فعال"
        if use_rag:
            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} ({arch})\n{msg_rag}"

    def answer(self, question: str, use_rag: bool, max_new_tokens: int, temperature: float, top_p: float, num_beams: int):
        if not question.strip():
            return "لطفاً سوال خود را وارد کنید.", ""
        if not self.gen:
            return "ابتدا مدل/RAG را بارگذاری کنید.", ""
        # update runtime params
        self.scfg.model.max_new_tokens = int(max_new_tokens)
        self.scfg.model.temperature = float(temperature)
        self.scfg.model.top_p = float(top_p)
        self.scfg.model.num_beams = int(num_beams)
        arts = self.rag.retrieve(question) if (use_rag and self.rag.collection) else []
        ctx = self.rag.build_context(arts) if arts else ""
        ans = self.gen.generate(question, ctx)
        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, arch: str, files: List[gr.File], use_rag: bool, epochs: int, batch: int, lr: float):
        self.scfg.model.model_name = model_name
        self.scfg.model.architecture = arch
        self.scfg.train.epochs = int(epochs)
        self.scfg.train.batch_size = int(batch)
        self.scfg.train.lr = float(lr)
        # ensure loader
        self.loader = ModelLoader(self.scfg.model).load()
        # train
        paths = [f.name for f in files] if files else []
        tm = TrainerManager(self.scfg, self.loader)
        if arch == "seq2seq":
            tm.train_seq2seq(paths, use_rag=use_rag)
        else:
            tm.train_causal(paths, use_rag=use_rag)
        return f"✅ آموزش کامل شد و در {self.scfg.train.output_dir} ذخیره شد."

    # --- UI ---
    def build_ui(self):
        default_models = {
            "Seq2Seq (mt5-base)": ("google/mt5-base", "seq2seq"),
            "Seq2Seq (t5-fa-base)": ("HooshvareLab/t5-fa-base", "seq2seq"),
            "Seq2Seq (flan-t5-base)": ("google/flan-t5-base", "seq2seq"),
            "Causal (Mistral-7B Instruct)": ("mistralai/Mistral-7B-Instruct-v0.2", "causal"),
        }
        with gr.Blocks(title="ماحون — مشاور حقوقی هوشمند", theme=gr.themes.Soft(primary_hue="green", secondary_hue="gray")) as app:
            gr.HTML("""
            <div style='text-align:center;padding:18px'>
              <h1 style='margin-bottom:4px'>ماحون — Ultimate Legal AI</h1>
              <p style='color:#666'>RAG • Seq2Seq/Causal • Training • Polished UI</p>
            </div>
            """)

            with gr.Tab("مشاوره"):
                with gr.Row():
                    model_dd = gr.Dropdown(choices=list(default_models.keys()), value="Seq2Seq (mt5-base)", label="مدل")
                    arch_info = gr.Markdown("""**راهنما:** مدل‌های Seq2Seq (MT5/T5) برای پاسخ‌های ساختاریافته عالی‌اند؛ مدل‌های Causal (Mistral) برای مکالمه طبیعی‌ترند.""")
                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, 10, 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("بارگذاری مدل/RAG", variant="primary")
                status = gr.Textbox(label="وضعیت", interactive=False)

                with gr.Accordion("پارامترهای تولید", open=False):
                    max_new_tokens = gr.Slider(64, 1024, 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")
                    num_beams = gr.Slider(1, 8, value=self.scfg.model.num_beams, step=1, label="num_beams (Seq2Seq)")

                question = gr.Textbox(lines=3, label="سوال حقوقی")
                examples = gr.Examples([
                    ["در صورت نقض قرارداد فروش، چه اقداماتی باید انجام دهم؟"],
                    ["آیا درج شرط عدم رقابت در قرارداد کار قانونی است؟"],
                    ["حق و حقوق کارگر در صورت اخراج فوری چیست؟"],
                    ["فرآیند طرح دعوای مطالبه مهریه چگونه است؟"],
                ], inputs=question, label="نمونه پرسش‌ها")
                ask_btn = gr.Button("پرسش", variant="primary")
                answer = gr.Markdown(label="پاسخ")
                refs = gr.Markdown(label="مواد قانونی مرتبط")

            with gr.Tab("آموزش"):
                gr.Markdown("برای آموزش، فایل‌های JSONL شامل کلیدهای `input` و `output` را بارگذاری کنید.")
                with gr.Row():
                    model_dd_train = gr.Dropdown(choices=list(default_models.keys()), value="Seq2Seq (mt5-base)", label="مدل")
                    use_rag_train = gr.Checkbox(value=True, label="RAG‑enhanced Training")
                train_files = gr.Files(label="JSONL Files", file_count="multiple", file_types=[".jsonl"])
                with gr.Row():
                    epochs = gr.Slider(1, 6, value=self.scfg.train.epochs, step=1, label="epochs")
                    batch = gr.Slider(1, 8, value=self.scfg.train.batch_size, step=1, label="batch per device")
                    lr = gr.Number(value=self.scfg.train.lr, label="learning rate")
                train_btn = gr.Button("شروع آموزش", variant="primary")
                train_status = gr.Textbox(label="وضعیت آموزش", interactive=False)

            # Events
            def _resolve(choice: str) -> Tuple[str,str]:
                return default_models[choice]

            load_btn.click(lambda choice, rag, pdir, coll, k, th: self.load(*_resolve(choice), rag, pdir, coll, k, th),
                           inputs=[model_dd, use_rag, persist_dir, collection, top_k, threshold], outputs=status)

            ask_btn.click(lambda q, rag, mnt, t, p, nb: self.answer(q, rag, mnt, t, p, nb),
                          inputs=[question, use_rag, max_new_tokens, temperature, top_p, num_beams], outputs=[answer, refs])

            train_btn.click(lambda choice, files, rag, e, b, l: self.train(*_resolve(choice), files, rag, e, b, l),
                            inputs=[model_dd_train, train_files, use_rag_train, epochs, batch, lr], outputs=train_status)
        return app

# ==========================

# Entrypoint
# ==========================
if __name__ == "__main__":
    app = LegalApp()
    ui = app.build_ui()
    ui.launch(share=True)