Update app.py

app.py

@@ -1,38 +1,35 @@
 # -*- coding: utf-8 -*-
 """
-Mahoon Legal AI — Enhanced Version
-Features:
-- Improved memory management and resource cleanup
-- Caching system for models and embeddings
-- Enhanced security and input validation
-- Better error handling and logging
-- Metrics and monitoring
-- Thread safety improvements
-- Configuration validation with Pydantic
-- Comprehensive testing support
-- Gradio UI with advanced features
+Mahoun — Legal AI (RAG + Training + Metrics) for HF Spaces / Gradio 5
+- سازگار با Gradio 5.x و Transformers >= 4.44
+- TrainingArguments ایمن با عقب‌سازگاری (safe_training_args)
+- RAG با ChromaDB + ایندکس‌سازی JSONL قوانین
+- متریک‌ها: ROUGE-L (seq2seq) و F1 ساده (causal)
+- ماسک پدینگ روی labels در معماری علّی
+- Progress به‌صورت DI: progress=gr.Progress(track_tqdm=True)
+
+ساختار ورودی دیتاست آموزش:
+  JSONL با کلیدهای "input" و "output"
+
+ساختار ورودی قوانین برای ایندکس:
+  JSONL با کلیدهای (پیش‌فرض) "article_id" و "text"
 """
 
 from __future__ import annotations
-import os
-import json
-import warnings
-import hashlib
-import threading
-import time
-from contextlib import contextmanager
+import os, sys, json, warnings
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import List, Dict, Optional, Tuple, Any, Union
-from datetime import datetime
-from functools import lru_cache
-import logging
+from typing import List, Dict, Optional, Tuple
 
+import numpy as np
 import torch
 from torch.utils.data import Dataset
 from sklearn.model_selection import train_test_split
-from pydantic import BaseModel, validator, Field
 
+import gradio as gr
+from packaging import version
+
+import transformers as tf
 from transformers import (
     AutoTokenizer,
     AutoModelForSeq2SeqLM,
@@ -41,1576 +38,684 @@ from transformers import (
     TrainingArguments,
     EarlyStoppingCallback,
     DataCollatorForSeq2Seq,
-    TrainerCallback
 )
 
+# RAG stack
 import chromadb
 from sentence_transformers import SentenceTransformer
-import gradio as gr
 
-warnings.filterwarnings("ignore")
+# Optional metrics
+try:
+    from evaluate import load as eval_load
+except Exception:
+    eval_load = None
 
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
+warnings.filterwarnings("ignore")
 
 # ==========================
-# Enhanced Config with Validation
+# Config
 # ==========================
-class ModelConfig(BaseModel):
-    model_name: str = "persiannlp/parsi-t5-base"
-    architecture: str = "seq2seq"
-    max_input_length: int = Field(default=1024, ge=64, le=4096)
-    max_target_length: int = Field(default=512, ge=32, le=2048)
-    max_new_tokens: int = Field(default=512, ge=32, le=1024)
-    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
-    top_p: float = Field(default=0.9, ge=0.1, le=1.0)
-    num_beams: int = Field(default=4, ge=1, le=8)
-    use_bf16: bool = True
-
-    @validator('architecture')
-    def validate_architecture(cls, v):
-        if v not in ['seq2seq', 'causal']:
-            raise ValueError('architecture must be seq2seq or causal')
-        return v
-
-    class Config:
-        validate_assignment = True
+@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
+    gradient_checkpointing: bool = True
 
-class SystemConfig(BaseModel):
-    model: ModelConfig = Field(default_factory=ModelConfig)
+@dataclass
+class RAGConfig:
     embedding_model: str = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
-    chroma_db_path: str = "./chroma.sqlite3"
-    top_k_retrieval: int = Field(default=5, ge=1, le=20)
-    similarity_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
-    cache_dir: str = "./cache"
-    output_dir: str = "./mahoon_legal_model"
+    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
+    enable: bool = True
+
+@dataclass
+class TrainConfig:
+    output_dir: str = "./mahoon_model"
     seed: int = 42
-    train_test_ratio: float = Field(default=0.1, ge=0.05, le=0.3)
-    batch_size: int = Field(default=2, ge=1, le=16)
-    grad_accum: int = Field(default=2, ge=1, le=8)
-    epochs: int = Field(default=2, ge=1, le=10)
-    lr: float = Field(default=3e-5, ge=1e-6, le=1e-3)
-    max_file_size_mb: int = Field(default=10, ge=1, le=100)
-    max_lines_per_file: int = Field(default=10000, ge=100, le=100000)
-    request_timeout: int = Field(default=30, ge=5, le=300)
-
-    class Config:
-        validate_assignment = True
+    test_size: float = 0.1
+    epochs: int = 3
+    batch_size: int = 2
+    grad_accum: int = 2
+    lr: float = 3e-5
+    use_bf16: bool = True
+    weight_decay: float = 0.01
+    warmup_ratio: float = 0.05
+    logging_steps: int = 50
+    eval_strategy: str = "epoch"   # "steps" | "epoch"
+    save_strategy: str = "epoch"
+    save_total_limit: int = 2
+    report_to: str = "none"        # "none" | "wandb"
+    max_grad_norm: float = 1.0
 
-# ==========================
-# Metrics and Monitoring
-# ==========================
 @dataclass
-class SystemMetrics:
-    requests_count: int = 0
-    avg_response_time: float = 0.0
-    error_count: int = 0
-    success_count: int = 0
-    memory_usage_mb: float = 0.0
-    last_updated: datetime = field(default_factory=datetime.now)
-    active_models: List[str] = field(default_factory=list)
-
-class MetricsCollector:
-    def __init__(self):
-        self.metrics = SystemMetrics()
-        self._lock = threading.Lock()
-
-    def record_request(self, response_time: float, success: bool = True):
-        with self._lock:
-            self.metrics.requests_count += 1
-            if success:
-                self.metrics.success_count += 1
-            else:
-                self.metrics.error_count += 1
-
-            # Update average response time
-            total_requests = self.metrics.requests_count
-            old_avg = self.metrics.avg_response_time
-            self.metrics.avg_response_time = (old_avg * (total_requests - 1) + response_time) / total_requests
-            self.metrics.last_updated = datetime.now()
-
-    def update_memory_usage(self):
-        if torch.cuda.is_available():
-            memory_mb = torch.cuda.memory_allocated() / 1024 / 1024
-            self.metrics.memory_usage_mb = memory_mb
-
-    def get_metrics(self) -> Dict[str, Any]:
-        with self._lock:
-            return {
-                "requests_total": self.metrics.requests_count,
-                "success_rate": self.metrics.success_count / max(self.metrics.requests_count, 1) * 100,
-                "avg_response_time": round(self.metrics.avg_response_time, 2),
-                "error_count": self.metrics.error_count,
-                "memory_usage_mb": round(self.metrics.memory_usage_mb, 2),
-                "active_models": self.metrics.active_models.copy(),
-                "last_updated": self.metrics.last_updated.isoformat()
-            }
-
-# Global metrics instance
-metrics = MetricsCollector()
+class SystemConfig:
+    model: ModelConfig = field(default_factory=ModelConfig)
+    rag: RAGConfig = field(default_factory=RAGConfig)
+    train: TrainConfig = field(default_factory=TrainConfig)
 
 # ==========================
-# Enhanced Utilities
+# Utils
 # ==========================
 def set_seed_all(seed: int = 42):
     import random
     random.seed(seed)
+    np.random.seed(seed)
     torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    logger.info(f"Set random seed to {seed}")
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
 
-def validate_file_security(file_path: str, max_size_mb: int = 10, max_lines: int = 10000) -> Tuple[bool, str]:
-    """Enhanced file validation with security checks"""
+def log_deps():
     try:
-        path = Path(file_path)
-
-        # Check if file exists and is readable
-        if not path.exists() or not path.is_file():
-            return False, "فایل وجود ندارد یا قابل خواندن نیست"
-
-        # Check file extension
-        if path.suffix.lower() != '.jsonl':
-            return False, "فقط فایل‌های .jsonl پذیرفته می‌شوند"
-
-        # Check file size
-        size_mb = path.stat().st_size / (1024 * 1024)
-        if size_mb > max_size_mb:
-            return False, f"حجم فایل نباید از {max_size_mb} مگابایت بیشتر باشد"
-
-        # Validate content structure
-        line_count = 0
-        with open(path, 'r', encoding='utf-8') as f:
-            for line_num, line in enumerate(f, 1):
-                line = line.strip()
-                if not line:
-                    continue
-
-                line_count += 1
-                if line_count > max_lines:
-                    return False, f"فایل نباید بیش از {max_lines} خط داشته باشد"
-
-                # Validate JSON structure
-                try:
-                    data = json.loads(line)
-                    if not isinstance(data, dict):
-                        return False, f"خط {line_num}: فرمت JSON نامعتبر"
-
-                    if 'input' not in data or 'output' not in data:
-                        return False, f"خط {line_num}: کلیدهای 'input' و 'output' الزامی هستند"
-
-                    # Check content length
-                    if len(str(data['input'])) > 2048 or len(str(data['output'])) > 2048:
-                        return False, f"خط {line_num}: طول محتوا بیش از حد مجاز"
-
-                except json.JSONDecodeError:
-                    return False, f"خط {line_num}: فرمت JSON نام��تبر"
-
-        if line_count == 0:
-            return False, "فایل خالی است"
-
-        return True, f"فایل معتبر است ({line_count} خط)"
-
+        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:
-        logger.error(f"File validation error: {e}")
-        return False, f"خطا در بررسی فایل: {str(e)}"
-
-def read_jsonl_files_safe(paths: List[str], cfg: SystemConfig) -> Tuple[List[Dict], List[str]]:
-    """Safe JSONL file reading with validation"""
-    data: List[Dict] = []
-    errors: List[str] = []
-
-    for path in paths:
-        # Validate file first
-        is_valid, msg = validate_file_security(path, cfg.max_file_size_mb, cfg.max_lines_per_file)
-        if not is_valid:
-            errors.append(f"{Path(path).name}: {msg}")
-            continue
-
-        try:
-            with open(path, 'r', encoding='utf-8') as f:
-                for line_num, line in enumerate(f, 1):
-                    line = line.strip()
-                    if not line:
-                        continue
-
-                    try:
-                        obj = json.loads(line)
-                        # Sanitize input
-                        obj['input'] = str(obj['input']).strip()
-                        obj['output'] = str(obj['output']).strip()
-
-                        if obj['input'] and obj['output']:
-                            data.append(obj)
-                    except json.JSONDecodeError:
-                        errors.append(f"{Path(path).name} line {line_num}: JSON decode error")
-
-        except Exception as e:
-            errors.append(f"{Path(path).name}: {str(e)}")
-
-    logger.info(f"Loaded {len(data)} samples from {len(paths)} files")
-    return data, errors
-
-# ==========================
-# Model Cache System
-# ==========================
-class ModelCache:
-    _instances: Dict[str, Any] = {}
-    _lock = threading.Lock()
-    _access_times: Dict[str, float] = {}
-    _max_cache_size = 3  # Maximum models to keep in cache
-
-    @classmethod
-    def _generate_key(cls, model_name: str, architecture: str) -> str:
-        return hashlib.md5(f"{model_name}_{architecture}".encode()).hexdigest()[:16]
-
-    @classmethod
-    def get_model(cls, model_name: str, architecture: str, model_config: ModelConfig):
-        key = cls._generate_key(model_name, architecture)
-
-        with cls._lock:
-            if key in cls._instances:
-                cls._access_times[key] = time.time()
-                logger.info(f"Model loaded from cache: {model_name}")
-                return cls._instances[key]
-
-            # Cleanup old models if cache is full
-            if len(cls._instances) >= cls._max_cache_size:
-                cls._cleanup_cache()
-
-            # Load new model
-            try:
-                loader = ModelLoader(model_config)
-                loader.load()
-                cls._instances[key] = loader
-                cls._access_times[key] = time.time()
-
-                # Update metrics
-                if model_name not in metrics.metrics.active_models:
-                    metrics.metrics.active_models.append(model_name)
-
-                logger.info(f"Model loaded and cached: {model_name}")
-                return loader
-
-            except Exception as e:
-                logger.error(f"Failed to load model {model_name}: {e}")
-                raise
-
-    @classmethod
-    def _cleanup_cache(cls):
-        """Remove least recently used model"""
-        if not cls._access_times:
-            return
-
-        # Find least recently used model
-        lru_key = min(cls._access_times.keys(), key=lambda k: cls._access_times[k])
-
-        # Clean up resources
-        if lru_key in cls._instances:
-            loader = cls._instances[lru_key]
-            cls._cleanup_model_resources(loader)
-            del cls._instances[lru_key]
-            del cls._access_times[lru_key]
-            logger.info(f"Removed model from cache: {lru_key}")
-
-    @classmethod
-    def _cleanup_model_resources(cls, loader):
-        """Clean up model resources"""
-        try:
-            if hasattr(loader, 'model') and hasattr(loader.model, 'cpu'):
-                loader.model.cpu()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-        except Exception as e:
-            logger.warning(f"Error cleaning up model resources: {e}")
-
-    @classmethod
-    def clear_cache(cls):
-        """Clear all cached models"""
-        with cls._lock:
-            for loader in cls._instances.values():
-                cls._cleanup_model_resources(loader)
-            cls._instances.clear()
-            cls._access_times.clear()
-            metrics.metrics.active_models.clear()
-            logger.info("Model cache cleared")
+        print("[deps] warn:", e, flush=True)
+
+def bf16_supported():
+    return torch.cuda.is_available() and hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported()
+
+def safe_training_args(**kwargs):
+    """
+    Wrapper برای سازگاری با نسخه‌های قدیمی‌تر Transformers (قبل از 4.4):
+    - evaluation_strategy -> evaluate_during_training
+    - حذف کلیدهای جدید که ممکن است ناشناخته باشند
+    """
+    tf_ver = version.parse(tf.__version__)
+    k = dict(kwargs)
+    if tf_ver < version.parse("4.4.0"):
+        eval_strat = k.pop("evaluation_strategy", None)
+        k["evaluate_during_training"] = bool(eval_strat and str(eval_strat).lower() != "no")
+        for rm in ["save_strategy","load_best_model_at_end","metric_for_best_model",
+                   "greater_is_better","predict_with_generate","generation_max_length",
+                   "generation_num_beams","report_to","max_grad_norm"]:
+            k.pop(rm, None)
+    return TrainingArguments(**k)
 
 # ==========================
-# Enhanced RAG System
+# RAG
 # ==========================
-class LegalRAGSystem:
-    def __init__(self, cfg: SystemConfig):
+class LegalRAG:
+    def __init__(self, cfg: RAGConfig):
         self.cfg = cfg
-        self.embedding_model: Optional[SentenceTransformer] = None
         self.client = None
         self.collection = None
-        self._lock = threading.Lock()
+        self.embedder: Optional[SentenceTransformer] = None
 
-    @contextmanager
-    def _safe_operation(self, operation_name: str):
-        """Context manager for safe RAG operations"""
-        start_time = time.time()
+    def init(self):
+        Path(self.cfg.persist_dir).mkdir(parents=True, exist_ok=True)
+        self.client = chromadb.PersistentClient(path=self.cfg.persist_dir)
         try:
-            yield
-        except Exception as e:
-            logger.error(f"RAG {operation_name} failed: {e}")
-            metrics.record_request(time.time() - start_time, success=False)
-            raise
-        else:
-            metrics.record_request(time.time() - start_time, success=True)
-
-    def setup_embedding(self):
-        if self.embedding_model is None:
-            try:
-                self.embedding_model = SentenceTransformer(
-                    self.cfg.embedding_model,
-                    cache_folder=self.cfg.cache_dir
-                )
-                logger.info(f"Embedding model loaded: {self.cfg.embedding_model}")
-            except Exception as e:
-                logger.error(f"Failed to load embedding model: {e}")
-                raise
-
-    def load_chroma(self) -> Tuple[bool, str]:
-        with self._safe_operation("load_chroma"):
+            self.collection = self.client.get_or_create_collection(self.cfg.collection)
+        except Exception:
             try:
-                base_path = str(Path(self.cfg.chroma_db_path).parent)
-                os.makedirs(base_path, exist_ok=True)
-
-                self.client = chromadb.PersistentClient(path=base_path)
+                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 index_jsonl(self, jsonl_path: str, id_key="article_id", text_key="text"):
+        """ایندکس‌سازی اولیه قوانین از JSONL: هر خط یک شیء {article_id, text, ...}."""
+        if not self.collection or not self.embedder:
+            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:
-                    self.collection = self.client.get_collection("legal_articles")
-                    count = self.collection.count()
-                    logger.info(f"Loaded existing collection with {count} documents")
-                    return count > 0, f"مجموعه موجود با {count} سند بارگذاری شد"
-                except Exception:
-                    self.collection = self.client.create_collection(
-                        "legal_articles",
-                        metadata={"description": "مواد قانونی"}
-                    )
-                    logger.info("Created new collection")
-                    return False, "مجموعه جدید ایجاد شد"
-
-            except Exception as e:
-                logger.error(f"ChromaDB initialization failed: {e}")
-                return False, f"خطا در بارگذاری پایگاه داده: {str(e)}"
+                    obj = json.loads(s)
+                except:
+                    continue
+                aid = str(obj.get(id_key, f"auto_{i}"))
+                txt = 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)
+        return f"✅ {len(ids)} سند قانونی ایندکس شد."
 
     def retrieve(self, query: str) -> List[Dict]:
-        if not self.collection or not query.strip():
+        if not self.collection:
             return []
-
-        with self._safe_operation("retrieve"):
-            try:
-                # Sanitize query
-                query = query.strip()[:500]  # Limit query length
-
-                result = self.collection.query(
-                    query_texts=[query],
-                    n_results=self.cfg.top_k_retrieval,
-                    include=["documents", "metadatas", "distances"]
-                )
-
-                articles = []
-                if result['documents'] and result['documents'][0]:
-                    for i, (doc, meta, dist) in enumerate(zip(
-                        result['documents'][0],
-                        result['metadatas'][0],
-                        result['distances'][0]
-                    )):
-                        similarity = max(0, min(1, 1 - dist))  # Normalize similarity
-                        if similarity >= self.cfg.similarity_threshold:
-                            articles.append({
-                                "article_id": meta.get("article_id", f"unknown_{i}"),
-                                "text": str(doc)[:500],  # Limit text length
-                                "similarity": round(similarity, 3),
-                            })
-
-                logger.info(f"Retrieved {len(articles)} relevant articles")
-                return articles
-
-            except Exception as e:
-                logger.error(f"Article retrieval failed: {e}")
-                return []
-
-    @staticmethod
-    def build_context(articles: List[Dict], limit_chars: int = 500) -> str:
-        if not articles:
-            return ""
-
-        context_parts = []
-        total_chars = 0
-
-        for article in articles:
-            text = article['text'][:limit_chars]
-            part = f"• ماده {article['article_id']}: {text}"
-
-            if total_chars + len(part) > limit_chars * 3:  # Max total context
-                break
-
-            context_parts.append(part)
-            total_chars += len(part)
-
-        return "مواد مرتبط:\n" + "\n".join(context_parts)
-
-# ==========================
-# Enhanced Formalizer
-# ==========================
-class Formalizer:
-    def __init__(self, model_name="erfan226/persian-t5-formality-transfer", device=None):
-        self.model_name = model_name
-        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.tokenizer = None
-        self.model = None
-        self._initialized = False
-        self._lock = threading.Lock()
-
-    def _initialize(self):
-        """Lazy initialization of formalizer model"""
-        if self._initialized:
-            return
-
-        with self._lock:
-            if self._initialized:  # Double-check pattern
-                return
-
-            try:
-                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, use_fast=True)
-                self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name).to(self.device)
-                self._initialized = True
-                logger.info("Formalizer model initialized")
-            except Exception as e:
-                logger.error(f"Formalizer initialization failed: {e}")
-                raise
-
-    def formalize(self, text: str, max_len: int = 512) -> str:
-        if not text or not text.strip():
-            return text
-
-        self._initialize()
-
         try:
-            # Sanitize and limit input
-            text = text.strip()[:max_len]
-
-            inputs = self.tokenizer(
-                text,
-                return_tensors="pt",
-                truncation=True,
-                max_length=max_len
-            ).to(self.device)
-
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_length=max_len,
-                    num_beams=4,
-                    early_stopping=True
-                )
-
-            result = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            logger.debug(f"Formalized text: {text[:50]}... -> {result[:50]}...")
-            return result
+            res = self.collection.query(
+                query_texts=[query],
+                n_results=self.cfg.top_k,
+                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)
+                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 []
 
-        except Exception as e:
-            logger.error(f"Text formalization failed: {e}")
-            return text  # Return original text on error
+    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)
 
 # ==========================
-# Enhanced Model Loader
+# Loader + Generator
 # ==========================
 class ModelLoader:
-    def __init__(self, model_config: ModelConfig):
-        self.cfg = model_config
+    def __init__(self, mcfg: ModelConfig):
+        self.cfg = mcfg
         self.tokenizer = None
         self.model = None
-        self._loaded = False
-
-    def _is_persianmind(self, name: str) -> bool:
-        return "PersianMind" in name or "universitytehran/PersianMind" in name
 
-    @contextmanager
-    def _gpu_memory_context(self):
-        """Context manager for GPU memory management"""
-        initial_memory = 0
+    def load(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(self.cfg.model_name)
+        # dtype انتخاب هوشمند
+        use_bf16 = bf16_supported() and self.cfg.gradient_checkpointing
+        dtype = torch.bfloat16 if use_bf16 else (torch.float16 if torch.cuda.is_available() else None)
+        model_kwargs = {"torch_dtype": dtype}
         if torch.cuda.is_available():
-            initial_memory = torch.cuda.memory_allocated()
+            model_kwargs["device_map"] = "auto"
+
+        if self.cfg.architecture == "seq2seq":
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(self.cfg.model_name, **model_kwargs)
+        elif self.cfg.architecture == "causal":
+            self.model = AutoModelForCausalLM.from_pretrained(self.cfg.model_name, **model_kwargs)
+            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")
 
-        try:
-            yield
-        finally:
-            if torch.cuda.is_available():
-                final_memory = torch.cuda.memory_allocated()
-                logger.info(f"Memory change: {(final_memory - initial_memory) / 1024**2:.1f} MB")
-                metrics.update_memory_usage()
-
-    def load(self, prefer_quantized: bool = True):
-        if self._loaded:
-            return self
-
-        with self._gpu_memory_context():
+        if self.cfg.gradient_checkpointing and hasattr(self.model, "gradient_checkpointing_enable"):
             try:
-                self._load_tokenizer()
-                self._load_model(prefer_quantized)
-                self._loaded = True
-                logger.info(f"Successfully loaded {self.cfg.model_name}")
-                return self
-
-            except Exception as e:
-                logger.error(f"Model loading failed: {e}")
-                self._cleanup()
-                raise
-
-    def _load_tokenizer(self):
-        """Load tokenizer with error handling"""
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(
-                self.cfg.model_name,
-                use_fast=True,
-                trust_remote_code=True
-            )
-            logger.info("Tokenizer loaded successfully")
-        except Exception as e:
-            logger.error(f"Tokenizer loading failed: {e}")
-            raise
-
-    def _load_model(self, prefer_quantized: bool):
-        """Load model with quantization support"""
-        device_map = "auto" if torch.cuda.is_available() else None
-        cuda_available = torch.cuda.is_available()
-        dtype = torch.bfloat16 if (cuda_available and self.cfg.use_bf16) else (
-            torch.float16 if cuda_available else torch.float32
-        )
-
-        # Try quantized loading for PersianMind causal models
-        if (self.cfg.architecture == "causal" and
-            self._is_persianmind(self.cfg.model_name) and
-            prefer_quantized and cuda_available):
-
-            if self._try_quantized_loading(device_map, dtype):
-                return
-
-        # Standard loading
-        self._load_standard_model(device_map, dtype)
-
-    def _try_quantized_loading(self, device_map, dtype) -> bool:
-        """Try loading model with quantization"""
-        # Try 8-bit first
-        try:
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.cfg.model_name,
-                device_map=device_map,
-                load_in_8bit=True,
-                torch_dtype=dtype,
-                trust_remote_code=True
-            )
-            self._setup_pad_token()
-            logger.info("Model loaded with 8-bit quantization")
-            return True
-        except Exception as e:
-            logger.warning(f"8-bit loading failed: {e}")
-
-        # Try 4-bit
-        try:
-            self.model = AutoModelForCausalLM.from_pretrained(
-                self.cfg.model_name,
-                device_map=device_map,
-                load_in_4bit=True,
-                bnb_4bit_use_double_quant=True,
-                bnb_4bit_compute_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16,
-                torch_dtype=dtype,
-                trust_remote_code=True
-            )
-            self._setup_pad_token()
-            logger.info("Model loaded with 4-bit quantization")
-            return True
-        except Exception as e:
-            logger.warning(f"4-bit loading failed: {e}")
-
-        return False
-
-    def _load_standard_model(self, device_map, dtype):
-        """Load model with standard precision"""
-        try:
-            if self.cfg.architecture == "seq2seq":
-                self.model = AutoModelForSeq2SeqLM.from_pretrained(
-                    self.cfg.model_name,
-                    device_map=device_map,
-                    torch_dtype=dtype,
-                    trust_remote_code=True
-                )
-            elif self.cfg.architecture == "causal":
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    self.cfg.model_name,
-                    device_map=device_map,
-                    torch_dtype=dtype,
-                    trust_remote_code=True
-                )
-                self._setup_pad_token()
-            else:
-                raise ValueError(f"Unsupported architecture: {self.cfg.architecture}")
-
-            logger.info("Model loaded with standard precision")
-
-        except Exception as e:
-            logger.error(f"Standard model loading failed: {e}")
-            raise
-
-    def _setup_pad_token(self):
-        """Setup pad token for causal models"""
-        if (self.tokenizer.pad_token is None and
-            hasattr(self.tokenizer, 'eos_token') and
-            self.tokenizer.eos_token):
-            self.tokenizer.pad_token = self.tokenizer.eos_token
-
-    def _cleanup(self):
-        """Clean up resources on failure"""
-        try:
-            if self.model and hasattr(self.model, 'cpu'):
-                self.model.cpu()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-        except Exception as e:
-            logger.warning(f"Cleanup error: {e}")
-
-# ==========================
-# Enhanced Generator
-# ==========================
-class UnifiedGenerator:
-    def __init__(self, loader: ModelLoader):
-        self.loader = loader
-        self.tokenizer = loader.tokenizer
+                self.model.gradient_checkpointing_enable()
+            except Exception:
+                pass
+        return self
+
+class Generator:
+    def __init__(self, loader: ModelLoader, mcfg: ModelConfig):
+        self.tk = loader.tokenizer
         self.model = loader.model
-        self.cfg = loader.cfg
-
-    def generate(self, question: str, context: str = "") -> Tuple[str, str]:
-        """Generate response with comprehensive error handling"""
-        if not question or not question.strip():
-            return "لطفاً سوال خود را وارد کنید.", "EMPTY_QUERY"
-
-        if not self.model or not self.tokenizer:
-            return "مدل بارگذاری نشده است.", "MODEL_NOT_LOADED"
-
-        start_time = time.time()
-        try:
-            # Sanitize inputs
-            question = question.strip()[:self.cfg.max_input_length // 2]
-            context = context.strip()[:self.cfg.max_input_length // 2]
-
-            if self.cfg.architecture == "seq2seq":
-                result = self._generate_seq2seq(question, context)
-            else:
-                result = self._generate_causal(question, context)
-
-            response_time = time.time() - start_time
-            metrics.record_request(response_time, success=True)
-
-            logger.info(f"Generated response in {response_time:.2f}s")
-            return result, ""
-
-        except torch.cuda.OutOfMemoryError:
-            error_msg = "حافظه GPU کافی نیست. لطفاً پارامترها را کاهش دهید."
-            logger.error("CUDA out of memory error")
-            metrics.record_request(time.time() - start_time, success=False)
-            return error_msg, "CUDA_OOM"
-
-        except Exception as e:
-            error_msg = "خطای غیرمنتظره در تولید پاسخ رخ داد."
-            logger.error(f"Generation error: {e}")
-            metrics.record_request(time.time() - start_time, success=False)
-            return error_msg, str(e)
-
-    def _generate_seq2seq(self, question: str, context: str) -> str:
-        """Generate response using seq2seq model"""
-        input_text = f"{context}\nسوال: {question}" if context else f"سوال: {question}"
-
-        inputs = self.tokenizer(
-            input_text,
-            return_tensors="pt",
-            truncation=True,
-            max_length=self.cfg.max_input_length
-        )
-        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
-
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
+        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,
-                no_repeat_ngram_size=2,
-                do_sample=False
             )
-
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        # Clean up response (remove input echo if present)
-        if input_text in response:
-            response = response.replace(input_text, "").strip()
-
-        return response or "پاسخی تولید نشد."
-
-    def _generate_causal(self, question: str, context: str) -> str:
-        """Generate response using causal model"""
-        prompt = f"{context}\nسوال: {question}\nپاسخ:" if context else f"سوال: {question}\nپاسخ:"
-
-        inputs = self.tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=self.cfg.max_input_length
-        )
-        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
-        input_length = inputs['input_ids'].shape[1]
-
-        with torch.no_grad():
-            outputs = self.model.generate(
-                **inputs,
+        else:
+            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=max(0.1, self.cfg.temperature),  # Ensure min temperature
+                temperature=self.cfg.temperature,
                 top_p=self.cfg.top_p,
-                pad_token_id=self.tokenizer.pad_token_id or self.tokenizer.eos_token_id,
-                repetition_penalty=1.1,
-                no_repeat_ngram_size=3
+                pad_token_id=self.tk.pad_token_id or self.tk.eos_token_id,
             )
-
-        # Extract only the generated part
-        generated_tokens = outputs[0][input_length:]
-        response = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
-
-        # Clean up response
-        response = response.strip()
-        if not response:
-            return "پاسخی تولید نشد."
-
-        # Remove any remaining prompt artifacts
-        response = response.split("سوال:")[0].strip()
-
-        return response
+        return self.tk.decode(out[0], skip_special_tokens=True)
 
 # ==========================
-# Enhanced Datasets
+# Datasets
 # ==========================
 class Seq2SeqJSONLDataset(Dataset):
-    def __init__(self, data: List[Dict], tokenizer, max_input: int, max_target: int):
-        self.tokenizer = tokenizer
-        self.max_input = max_input
-        self.max_target = max_target
-
-        # Filter and validate data
+    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 item in data:
-            src = str(item.get("input", "")).strip()
-            tgt = str(item.get("output", "")).strip()
-
-            if src and tgt and len(src) > 5 and len(tgt) > 5:  # Minimum length check
-                self.items.append((src, tgt))
-
-        logger.info(f"Seq2Seq dataset created with {len(self.items)} samples")
+        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):
-        source_text, target_text = self.items[idx]
-
-        # Tokenize inputs
-        model_inputs = self.tokenizer(
-            source_text,
-            max_length=self.max_input,
-            padding="max_length",
-            truncation=True,
-            return_tensors="pt"
-        )
-
-        # Tokenize targets
-        labels = self.tokenizer(
-            text_target=target_text,
-            max_length=self.max_target,
-            padding="max_length",
-            truncation=True,
-            return_tensors="pt"
-        )
-
-        # Convert to proper format
-        return {
-            "input_ids": model_inputs["input_ids"].flatten(),
-            "attention_mask": model_inputs["attention_mask"].flatten(),
-            "labels": labels["input_ids"].flatten()
-        }
+        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_length: int):
-        self.tokenizer = tokenizer
-        self.max_length = max_length
-
-        # Process data
+    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 item in data:
-            src = str(item.get("input", "")).strip()
-            tgt = str(item.get("output", "")).strip()
-
-            if src and tgt and len(src) > 5 and len(tgt) > 5:
-                formatted_text = f"سوال: {src}\nپاسخ: {tgt}"
-                self.items.append(formatted_text)
-
-        logger.info(f"Causal dataset created with {len(self.items)} samples")
+        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]
-
-        encoding = self.tokenizer(
-            text,
-            max_length=self.max_length,
-            padding="max_length",
-            truncation=True,
-            return_tensors="pt"
-        )
-
-        input_ids = encoding["input_ids"].flatten()
-        attention_mask = encoding["attention_mask"].flatten()
-
+        enc = self.tk(text, max_length=self.max_inp, padding="max_length", truncation=True)
+        input_ids = torch.tensor(enc["input_ids"])
+        attn = torch.tensor(enc["attention_mask"])
         labels = input_ids.clone()
-        labels[attention_mask == 0] = -100
-
-        return {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "labels": labels
-        }
+        labels[attn == 0] = -100  # padding mask for loss
+        return {"input_ids": input_ids, "attention_mask": attn, "labels": labels}
 
 # ==========================
-# Enhanced Progress Callback
+# Metrics
 # ==========================
-class GradioProgressCallback(TrainerCallback):
-    def __init__(self, progress: gr.Progress, status_textbox: gr.Textbox):
-        self.progress = progress
-        self.status_textbox = status_textbox
-        self.total_steps = None
-        self.start_time = None
-        self.last_update = 0
-
-    def on_train_begin(self, args, state, control, **kwargs):
-        self.total_steps = state.max_steps
-        self.start_time = time.time()
-        self.progress(0, desc="آموزش شروع شد 🚀")
-        self.status_textbox.update(value="آموزش شروع شد...")
-
-    def on_step_end(self, args, state, control, **kwargs):
-        if not self.total_steps or time.time() - self.last_update < 1.0:  # Throttle updates
-            return
-
-        self.last_update = time.time()
-
-        # Calculate progress
-        progress_pct = min(100, int((state.global_step / self.total_steps) * 100))
-
-        # Estimate remaining time
-        elapsed = time.time() - self.start_time
-        if state.global_step > 0:
-            avg_time_per_step = elapsed / state.global_step
-            remaining_steps = self.total_steps - state.global_step
-            eta_seconds = avg_time_per_step * remaining_steps
-            eta_minutes = int(eta_seconds / 60)
-            eta_str = f" (تخمین باقی‌مانده: {eta_minutes} دقیقه)" if eta_minutes > 0 else ""
-        else:
-            eta_str = ""
-
-        # Update progress
-        self.progress(progress_pct, desc=f"آموزش: {progress_pct}%")
-
-        # Update status with more details
-        current_lr = state.learning_rate if hasattr(state, 'learning_rate') else args.learning_rate
-        status_msg = (f"Step {state.global_step}/{self.total_steps} → {progress_pct}%{eta_str}\n"
-                     f"Learning Rate: {current_lr:.2e}")
-
-        if hasattr(state, 'log_history') and state.log_history:
-            last_log = state.log_history[-1]
-            if 'train_loss' in last_log:
-                status_msg += f"\nTrain Loss: {last_log['train_loss']:.4f}"
-            if 'eval_loss' in last_log:
-                status_msg += f"\nEval Loss: {last_log['eval_loss']:.4f}"
-
-        self.status_textbox.update(value=status_msg)
-
-    def on_evaluate(self, args, state, control, **kwargs):
-        if hasattr(state, 'log_history') and state.log_history:
-            last_log = state.log_history[-1]
-            if 'eval_loss' in last_log:
-                self.status_textbox.update(
-                    value=f"ارزیابی انجام شد - Loss: {last_log['eval_loss']:.4f}"
-                )
-
-    def on_train_end(self, args, state, control, **kwargs):
-        total_time = time.time() - self.start_time
-        total_minutes = int(total_time / 60)
-
-        self.progress(100, desc="آموزش تکمیل شد ✅")
-        self.status_textbox.update(
-            value=f"آموزش با موفقیت تکمیل شد ✅\n"
-                  f"زمان کل: {total_minutes} دقیقه\n"
-                  f"کل Steps: {state.global_step}"
-        )
+def build_metrics_fn(arch: str, tokenizer):
+    rouge = eval_load("rouge") if eval_load else None
+
+    def _postprocess(preds):
+        if isinstance(preds, (list, tuple)):
+            return [p.strip() for p in preds]
+        return preds
+
+    def compute_metrics_seq2seq(eval_pred):
+        if rouge is None:
+            return {"rougeL": 0.0}
+        preds, labels = eval_pred
+        if isinstance(preds, tuple):
+            preds = preds[0]
+        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+        decoded_preds = _postprocess(decoded_preds)
+        decoded_labels = _postprocess(decoded_labels)
+        r = rouge.compute(predictions=decoded_preds, references=decoded_labels, rouge_types=["rougeL"])
+        return {"rougeL": float(r.get("rougeL", 0.0))}
+
+    def compute_metrics_causal(eval_pred):
+        preds, labels = eval_pred
+        if isinstance(preds, tuple):
+            preds = preds[0]
+        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+        tp = fp = fn = 0
+        for p, g in zip(decoded_preds, decoded_labels):
+            p_set, g_set = set(p.split()), set(g.split())
+            tp += len(p_set & g_set)
+            fp += len(p_set - g_set)
+            fn += len(g_set - p_set)
+        precision = tp / (tp + fp + 1e-8)
+        recall = tp / (tp + fn + 1e-8)
+        f1 = 2 * precision * recall / (precision + recall + 1e-8)
+        return {"f1_simple": float(f1)}
+
+    return compute_metrics_seq2seq if arch == "seq2seq" else compute_metrics_causal
 
 # ==========================
-# Enhanced Trainer Manager
+# Trainer Manager
 # ==========================
-class TrainerManager:
-    def __init__(self, system_config: SystemConfig, model_loader: ModelLoader):
-        self.cfg = system_config
-        self.loader = model_loader
-
-    def train(self, train_paths: List[str], extra_callbacks: List = None) -> Tuple[bool, str]:
-        """Main training method with comprehensive error handling"""
-        if extra_callbacks is None:
-            extra_callbacks = []
-
-        try:
-            # Validate training files
-            data, errors = read_jsonl_files_safe(train_paths, self.cfg)
-
-            if errors:
-                error_msg = "خطاهای فایل:\n" + "\n".join(errors[:5])  # Show first 5 errors
-                return False, error_msg
-
-            if len(data) < 10:
-                return False, f"داده کافی نیست. حداقل 10 نمونه نیاز است (موجود: {len(data)})"
-
-            # Set random seed
-            set_seed_all(self.cfg.seed)
-
-            # Split data
-            train_data, val_data = train_test_split(
-                data,
-                test_size=self.cfg.train_test_ratio,
-                random_state=self.cfg.seed
-            )
-
-            logger.info(f"Training samples: {len(train_data)}, Validation samples: {len(val_data)}")
-
-            # Train based on architecture
-            if self.cfg.model.architecture == "seq2seq":
-                success, msg = self._train_seq2seq(train_data, val_data, extra_callbacks)
-            else:
-                success, msg = self._train_causal(train_data, val_data, extra_callbacks)
-
-            if success:
-                # Save configuration
-                self._save_training_config()
-
-            return success, msg
-
-        except Exception as e:
-            logger.error(f"Training failed: {e}")
-            return False, f"خطا در آموزش: {str(e)}"
-
-    def _train_seq2seq(self, train_data: List[Dict], val_data: List[Dict],
-                      extra_callbacks: List) -> Tuple[bool, str]:
-        """Train seq2seq model"""
-        try:
-            # Create datasets
-            train_dataset = Seq2SeqJSONLDataset(
-                train_data, self.loader.tokenizer,
-                self.cfg.model.max_input_length,
-                self.cfg.model.max_target_length
-            )
-
-            val_dataset = Seq2SeqJSONLDataset(
-                val_data, self.loader.tokenizer,
-                self.cfg.model.max_input_length,
-                self.cfg.model.max_target_length
-            )
-
-            # Data collator
-            data_collator = DataCollatorForSeq2Seq(
-                tokenizer=self.loader.tokenizer,
-                model=self.loader.model,
-                padding=True
-            )
-
-            # Training arguments
-            training_args = self._get_training_args()
-            training_args.predict_with_generate = True
-            training_args.generation_max_length = self.cfg.model.max_target_length
-            training_args.generation_num_beams = self.cfg.model.num_beams
-
-            # Create trainer
-            trainer = Trainer(
-                model=self.loader.model,
-                args=training_args,
-                train_dataset=train_dataset,
-                eval_dataset=val_dataset,
-                data_collator=data_collator,
-                tokenizer=self.loader.tokenizer,
-                callbacks=self._get_callbacks(extra_callbacks)
-            )
-
-            # Train
-            trainer.train()
-
-            # Save model
-            trainer.save_model(self.cfg.output_dir)
-            self.loader.tokenizer.save_pretrained(self.cfg.output_dir)
-
-            return True, "مدل Seq2Seq با موفقیت آموزش داده شد"
-
-        except Exception as e:
-            logger.error(f"Seq2Seq training failed: {e}")
-            return False, f"خطا در آموزش Seq2Seq: {str(e)}"
-
-    def _train_causal(self, train_data: List[Dict], val_data: List[Dict],
-                     extra_callbacks: List) -> Tuple[bool, str]:
-        """Train causal language model"""
-        try:
-            # Create datasets
-            train_dataset = CausalJSONLDataset(
-                train_data, self.loader.tokenizer,
-                self.cfg.model.max_input_length
-            )
-
-            val_dataset = CausalJSONLDataset(
-                val_data, self.loader.tokenizer,
-                self.cfg.model.max_input_length
-            )
-
-            # Training arguments
-            training_args = self._get_training_args()
-
-            # Create trainer
-            trainer = Trainer(
-                model=self.loader.model,
-                args=training_args,
-                train_dataset=train_dataset,
-                eval_dataset=val_dataset,
-                tokenizer=self.loader.tokenizer,
-                callbacks=self._get_callbacks(extra_callbacks)
-            )
-
-            # Train
-            trainer.train()
-
-            # Save model
-            trainer.save_model(self.cfg.output_dir)
-            self.loader.tokenizer.save_pretrained(self.cfg.output_dir)
+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
 
-            return True, "مدل Causal با موفقیت آموزش داده شد"
+class TrainerManager:
+    def __init__(self, syscfg: SystemConfig, loader: ModelLoader):
+        self.cfg = syscfg
+        self.loader = loader
 
-        except Exception as e:
-            logger.error(f"Causal training failed: {e}")
-            return False, f"خطا در آموزش Causal: {str(e)}"
-
-    def _get_training_args(self) -> TrainingArguments:
-        """Get training arguments with optimized settings"""
-        return TrainingArguments(
-            output_dir=self.cfg.output_dir,
-            num_train_epochs=self.cfg.epochs,
-            learning_rate=self.cfg.lr,
-            per_device_train_batch_size=self.cfg.batch_size,
-            per_device_eval_batch_size=self.cfg.batch_size,
-            gradient_accumulation_steps=self.cfg.grad_accum,
-            warmup_ratio=0.05,
-            weight_decay=0.01,
-            evaluation_strategy="epoch",
-            eval_steps=500,
-            save_strategy="epoch",
-            save_total_limit=3,  # Keep more checkpoints
+    def _args_common(self, is_seq2seq: bool):
+        fp16_ok = torch.cuda.is_available() and (not self.cfg.train.use_bf16)
+        bf16_ok = bf16_supported() and self.cfg.train.use_bf16
+
+        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",
-            greater_is_better=False,
-            logging_steps=50,
-            logging_dir=f"{self.cfg.output_dir}/logs",
-            report_to="none",
-            bf16=self.cfg.model.use_bf16 if torch.cuda.is_available() else False,
-            fp16=(not self.cfg.model.use_bf16) if torch.cuda.is_available() else False,
-            dataloader_drop_last=True,
-            remove_unused_columns=False,
-            gradient_checkpointing=True,  # Save memory
+            logging_steps=self.cfg.train.logging_steps,
+            report_to=([] if self.cfg.train.report_to == "none" else [self.cfg.train.report_to]),
+            fp16=fp16_ok,
+            bf16=bf16_ok,
+            max_grad_norm=self.cfg.train.max_grad_norm,
+            **({
+                "predict_with_generate": True,
+                "generation_max_length": self.cfg.model.max_target_length,
+                "generation_num_beams": self.cfg.model.num_beams
+            } if is_seq2seq else {})
         )
-
-    def _get_callbacks(self, extra_callbacks: List) -> List:
-        """Get training callbacks"""
-        callbacks = [
-            EarlyStoppingCallback(early_stopping_patience=3, early_stopping_threshold=0.01)
-        ]
-        callbacks.extend(extra_callbacks)
-        return callbacks
-
-    def _save_training_config(self):
-        """Save training configuration"""
-        try:
-            config_path = Path(self.cfg.output_dir) / "training_config.json"
-            config_dict = self.cfg.dict()
-            config_dict['training_timestamp'] = datetime.now().isoformat()
-            config_dict['training_completed'] = True
-
-            with open(config_path, 'w', encoding='utf-8') as f:
-                json.dump(config_dict, f, ensure_ascii=False, indent=2)
-
-            logger.info(f"Training config saved to {config_path}")
-        except Exception as e:
-            logger.warning(f"Failed to save training config: {e}")
+        return args
+
+    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 and self.cfg.rag.enable) 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)
+
+        args = self._args_common(is_seq2seq=True)
+        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)],
+            compute_metrics=build_metrics_fn("seq2seq", self.loader.tokenizer)
+        )
+        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 and self.cfg.rag.enable) 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)
+
+        args = self._args_common(is_seq2seq=False)
+        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)],
+            compute_metrics=build_metrics_fn("causal", self.loader.tokenizer)
+        )
+        trainer.train()
+        trainer.save_model(self.cfg.train.output_dir)
+        self.loader.tokenizer.save_pretrained(self.cfg.train.output_dir)
 
 # ==========================
-# Enhanced Legal App
+# App (Gradio 5)
 # ==========================
 class LegalApp:
-    def __init__(self, system_config: Optional[SystemConfig] = None):
-        self.cfg = system_config or SystemConfig()
-        self.rag = LegalRAGSystem(self.cfg)
-        self.formalizer: Optional[Formalizer] = None
-        self._current_loader: Optional[ModelLoader] = None
-        self._current_generator: Optional[UnifiedGenerator] = None
-        self._lock = threading.Lock()
-
-    def _ensure_model(self, model_name: str, architecture: str) -> Tuple[bool, str]:
-        """Ensure model is loaded with thread safety"""
-        with self._lock:
+    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
+
+    # --- helpers ---
+    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 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)
+        self.scfg.rag.enable = bool(use_rag)
+
+        # 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:
-                # Update config
-                self.cfg.model.model_name = model_name
-                self.cfg.model.architecture = architecture
-
-                # Get model from cache
-                self._current_loader = ModelCache.get_model(model_name, architecture, self.cfg.model)
-                self._current_generator = UnifiedGenerator(self._current_loader)
-
-                return True, f"مدل بارگذاری شد: {model_name} ({architecture})"
-
+                self.rag = LegalRAG(self.scfg.rag)
+                self.rag.init()
+                msg_rag = "RAG آماده است"
             except Exception as e:
-                logger.error(f"Model loading failed: {e}")
-                return False, f"خطا در بارگذاری مدل: {str(e)}"
-
-    def _ensure_rag(self) -> Tuple[bool, str]:
-        """Ensure RAG system is ready"""
-        try:
-            self.rag.setup_embedding()
-            success, message = self.rag.load_chroma()
-            return success, message
-        except Exception as e:
-            logger.error(f"RAG setup failed: {e}")
-            return False, f"خطا در راه‌اندازی RAG: {str(e)}"
-
-    def _ensure_formalizer(self) -> str:
-        """Ensure formalizer is ready"""
-        try:
-            if not self.formalizer:
-                self.formalizer = Formalizer()
-            return "Formalizer آماده است."
-        except Exception as e:
-            logger.error(f"Formalizer setup failed: {e}")
-            return f"خطا در راه‌اندازی Formalizer: {str(e)}"
-
-    # Event handlers
-    def handle_load_model(self, model_choice: str, use_rag: bool) -> str:
-        """Handle model loading"""
-        try:
-            model_configs = self._get_model_configs()
-            if model_choice not in model_configs:
-                return "مدل نامعتبر انتخاب شده"
+                msg_rag = f"RAG خطا: {e}"
 
-            model_name, architecture = model_configs[model_choice]
-
-            # Load model
-            success, model_msg = self._ensure_model(model_name, architecture)
-            if not success:
-                return model_msg
-
-            # Setup RAG if requested
-            rag_msg = ""
-            if use_rag:
-                rag_success, rag_msg = self._ensure_rag()
-                rag_msg = f"\nRAG: {rag_msg}"
-            else:
-                rag_msg = "\nRAG: غیر فعال"
-
-            return f"{model_msg}{rag_msg}"
-
-        except Exception as e:
-            logger.error(f"Model loading handler failed: {e}")
-            return f"خطا در بارگذاری: {str(e)}"
-
-    def handle_generate_response(self, question: str, use_rag: bool, use_formalizer: bool,
-                                max_new_tokens: int, temperature: float, top_p: float,
-                                num_beams: int) -> Tuple[str, str, str]:  # response, references, metrics
-        """Handle response generation"""
-        if not question or not question.strip():
-            return "لطفاً سوال خود را وارد کنید.", "", ""
-
-        if not self._current_generator:
-            return "ابتدا مدل را بارگذاری کنید.", "", ""
-
-        start_time = time.time()
+        return f"مدل بارگذاری شد: {model_name} ({arch})\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:
-            # Update generation parameters
-            self.cfg.model.max_new_tokens = max(32, min(1024, int(max_new_tokens)))
-            self.cfg.model.temperature = max(0.1, min(2.0, float(temperature)))
-            self.cfg.model.top_p = max(0.1, min(1.0, float(top_p)))
-            self.cfg.model.num_beams = max(1, min(8, int(num_beams)))
-
-            # Apply input formalization if requested
-            processed_question = question
-            if use_formalizer:
-                formalizer_msg = self._ensure_formalizer()
-                if "خطا" not in formalizer_msg and self.formalizer:
-                    processed_question = self.formalizer.formalize(question)
-
-            # Retrieve relevant articles if RAG is enabled
-            articles = []
-            if use_rag and self.rag.collection:
-                articles = self.rag.retrieve(processed_question)
-
-            # Build context
-            context = LegalRAGSystem.build_context(articles) if articles else ""
-
-            # Generate response
-            response, error = self._current_generator.generate(processed_question, context)
-
-            # Build references section
-            references = ""
-            if articles:
-                ref_parts = []
-                for article in articles[:3]:  # Limit to top 3 references
-                    ref_parts.append(
-                        f"**ماده {article['article_id']}** (شباهت: {article['similarity']:.2f})\n"
-                        f"{article['text'][:400]}{'...' if len(article['text']) > 400 else ''}"
-                    )
-                references = "\n\n".join(ref_parts)
-
-            # Generate metrics info
-            elapsed_time = time.time() - start_time
-            metrics_info = f"زمان پردازش: {elapsed_time:.2f}s"
-            if articles:
-                metrics_info += f" | مواد یافت شده: {len(articles)}"
-            if use_formalizer:
-                metrics_info += " | فرمالایزر فعال"
-
-            return response, references, metrics_info
-
+            self.rag.init()
+            p = getattr(laws_file, "name", None) or getattr(laws_file, "path", None)
+            if not p:
+                return "فایل قوانین معتبر نیست."
+            res = self.rag.index_jsonl(p, id_key=id_key, text_key=text_key)
+            return res
         except Exception as e:
-            logger.error(f"Response generation failed: {e}")
-            error_time = time.time() - start_time
-            metrics.record_request(error_time, success=False)
-            return f"خطا در تولید پاسخ: {str(e)}", "", f"خطا پس از {error_time:.2f}s"
-
-    def handle_training(self, model_choice: str, uploaded_files, use_rag_training: bool,
-                       epochs: int, batch_size: int, learning_rate: float,
-                       progress: gr.Progress, status_textbox: gr.Textbox) -> str:
-        """Handle model training"""
-        try:
-            # Validate inputs
-            if not uploaded_files:
-                return "لطفاً فایل‌های آموزشی را بارگذاری کنید."
-
-            # Get model config
-            model_configs = self._get_model_configs()
-            if model_choice not in model_configs:
-                return "مدل نامعتبر انتخاب شده"
-
-            model_name, architecture = model_configs[model_choice]
-
-            # Load model for training
-            success, msg = self._ensure_model(model_name, architecture)
-            if not success:
-                return f"خطا در بارگذاری مدل: {msg}"
-
-            # Update training config
-            self.cfg.epochs = max(1, min(10, int(epochs)))
-            self.cfg.batch_size = max(1, min(16, int(batch_size)))
-            self.cfg.lr = max(1e-6, min(1e-3, float(learning_rate)))
-
-            # Setup RAG if requested
-            if use_rag_training:
-                rag_success, rag_msg = self._ensure_rag()
-                if not rag_success:
-                    logger.warning(f"RAG setup failed for training: {rag_msg}")
-
-            # Get file paths (gr.File with type="filepath" returns list[str])
-            file_paths = uploaded_files
-
-            if not file_paths:
-                return "فایل‌های معتبر یافت نشد."
-
-            # Create trainer
-            trainer_manager = TrainerManager(self.cfg, self._current_loader)
-
-            # Create progress callback
-            progress_callback = GradioProgressCallback(progress, status_textbox)
-
-            # Start training
-            success, result_msg = trainer_manager.train(file_paths, [progress_callback])
-
-            if success:
-                # Clear model cache to force reload of trained model
-                ModelCache.clear_cache()
-                return f"✅ {result_msg}\nمدل در مسیر '{self.cfg.output_dir}' ذخیره شد."
-            else:
-                return f"❌ {result_msg}"
-
-        except Exception as e:
-            logger.error(f"Training handler failed: {e}")
-            return f"خطا در آموزش: {str(e)}"
-
-    def get_system_status(self) -> str:
-        """Get system status information"""
-        try:
-            status_parts = []
-
-            # Model status
-            if self._current_loader:
-                status_parts.append(f"✅ مدل فعال: {self.cfg.model.model_name}")
-            else:
-                status_parts.append("❌ مدل بارگذاری نشده")
-
-            # RAG status
-            if self.rag.collection:
-                doc_count = self.rag.collection.count()
-                status_parts.append(f"✅ RAG فعال ({doc_count} سند)")
-            else:
-                status_parts.append("❌ RAG غیر فعال")
-
-            # System metrics
-            sys_metrics = metrics.get_metrics()
-            status_parts.append(f"📊 درخواست‌ها: {sys_metrics['requests_total']}")
-            status_parts.append(f"📈 نرخ موفقیت: {sys_metrics['success_rate']:.1f}%")
-            status_parts.append(f"⏱️ زمان متوسط: {sys_metrics['avg_response_time']}s")
-
-            if torch.cuda.is_available():
-                memory_mb = torch.cuda.memory_allocated() / 1024 / 1024
-                status_parts.append(f"🖥️ حافظه GPU: {memory_mb:.1f} MB")
-
-            return "\n".join(status_parts)
+            return f"خطا در ایندکس: {e}"
+
+    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 را بارگذاری کنید.", ""
+        # 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.scfg.rag.enable 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,
+              wd: float, warmup: float, report_to: str, progress=gr.Progress(track_tqdm=True)):
+        progress(0.0, desc="راه‌اندازی")
+        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)
+        self.scfg.train.weight_decay = float(wd)
+        self.scfg.train.warmup_ratio = float(warmup)
+        self.scfg.train.report_to = report_to
+
+        progress(0.1, desc="بارگذاری مدل/توکنایزر")
+        self.loader = ModelLoader(self.scfg.model).load()
+
+        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.3, desc="آماده‌سازی دیتاست‌ها و RAG")
+        if arch == "seq2seq":
+            tm.train_seq2seq(paths, use_rag=use_rag)
+        else:
+            tm.train_causal(paths, use_rag=use_rag)
 
-        except Exception as e:
-            return f"خطا در دریافت وضعیت: {str(e)}"
+        progress(0.95, desc="ذخیرهٔ آرتیفکت‌ها")
+        return f"✅ آموزش کامل شد و در {self.scfg.train.output_dir} ذخیره شد."
 
-    def _get_model_configs(self) -> Dict[str, Tuple[str, str]]:
-        """Get available model configurations"""
-        return {
-            "Seq2Seq (parsi-t5-base)": ("persiannlp/parsi-t5-base", "seq2seq"),
+    # --- UI ---
+    def build_ui(self):
+        log_deps()
+        default_models = {
             "Seq2Seq (mt5-base)": ("google/mt5-base", "seq2seq"),
-            "Causal (Mistral-7B)": ("mistralai/Mistral-7B-Instruct-v0.2", "causal"),
-            "Causal (PersianMind-v1.0)": ("universitytehran/PersianMind-v1.0", "causal"),
+            "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"),
         }
 
-    def build_ui(self) -> gr.Blocks:
-        """Build enhanced Gradio interface"""
-        model_choices = list(self._get_model_configs().keys())
-
-        with gr.Blocks(
-            title="ماحون — مشاور حقوقی هوشمند",
-            theme=gr.themes.Soft(),
-            css="""
-            .status-box { font-family: 'Courier New', monospace; font-size: 12px; }
-            .metrics-box { background-color: #f0f0f0; padding: 10px; border-radius: 5px; }
-            """
-        ) as app:
-
+        with gr.Blocks(title="ماحون — مشاور حقوقی هوشمند", theme=gr.themes.Soft(primary_hue="green", secondary_hue="gray")) as app:
             gr.HTML("""
-            <div style='text-align: center; margin-bottom: 20px;'>
-                <h1>ماحون — مشاور حقوقی هوشمند 🏛️</h1>
-                <p>سیستم پیشرفته مشاوره حقوقی با قابلیت RAG، Fine-tuning و هوش مصنوعی</p>
+            <div style='text-align:center;padding:18px'>
+              <h1 style='margin-bottom:4px'>ماحون — Ultimate Legal AI</h1>
+              <p style='color:#666'>RAG • Seq2Seq/Causal • Training • Metrics</p>
             </div>
             """)
 
-            # System Status
-            with gr.Accordion("وضعیت سیستم", open=False):
-                system_status = gr.Markdown(
-                    value=self.get_system_status(),
-                    elem_classes=["status-box"]
+            with gr.Tab("مشاوره"):
+                with gr.Row():
+                    model_dd = gr.Dropdown(choices=list(default_models.keys()), value="Seq2Seq (mt5-base)", label="مدل")
+                    gr.Markdown("**راهنما:** Seq2Seq برای پاسخ‌های ساختاریافته؛ Causal برای مکالمه طبیعی‌تر.")
+                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("بارگذاری مدل/RAG", variant="primary")
+                status = gr.Textbox(label="وضعیت", interactive=False)
+
+                with gr.Accordion("ساخت ایندکس قوانین (اختیاری)", open=False):
+                    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)
+
+                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="سوال حقوقی")
+                gr.Examples(
+                    examples=[
+                        ["در صورت نقض قرارداد فروش، چه اقداماتی باید انجام دهم؟"],
+                        ["آیا درج شرط عدم رقابت در قرارداد کار قانونی است؟"],
+                        ["حق و حقوق کارگر در صورت اخراج فوری چیست؟"],
+                        ["فرآیند طرح دعوای مطالبه مهریه چگونه است؟"],
+                    ],
+                    inputs=question, label="نمونه پرسش‌ها"
                 )
-                refresh_status_btn = gr.Button("🔄 بروزرسانی وضعیت", size="sm")
-
-            with gr.Tabs() as tabs:
-                # Consultation Tab
-                with gr.Tab("💬 مشاوره") as advice_tab:
-                    with gr.Row():
-                        with gr.Column(scale=2):
-                            model_dropdown = gr.Dropdown(
-                                choices=model_choices,
-                                value=model_choices[0],
-                                label="انتخاب مدل",
-                                info="نوع مدل مورد نظر را انتخاب کنید"
-                            )
-                        with gr.Column(scale=1):
-                            use_rag_checkbox = gr.Checkbox(
-                                value=True,
-                                label="استفاده از RAG",
-                                info="بازیابی مواد قانونی مرتبط"
-                            )
-                            use_formalizer_checkbox = gr.Checkbox(
-                                value=False,
-                                label="رسمی‌سازی ورودی",
-                                info="تبدیل متن غیررسمی به رسمی"
-                            )
-
-                    load_model_btn = gr.Button("🚀 بارگذاری مدل/RAG", variant="primary", size="lg")
-                    load_status = gr.Textbox(
-                        label="وضعیت بارگذاری",
-                        interactive=False,
-                        elem_classes=["status-box"]
-                    )
-
-                    # Generation Parameters
-                    with gr.Accordion("⚙️ پارامترهای تولید", open=False):
-                        with gr.Row():
-                            max_new_tokens = gr.Slider(
-                                minimum=32, maximum=1024, value=self.cfg.model.max_new_tokens,
-                                step=16, label="حداکثر توکن‌های جدید"
-                            )
-                            temperature = gr.Slider(
-                                minimum=0.1, maximum=2.0, value=self.cfg.model.temperature,
-                                step=0.05, label="دما (خلاقیت)"
-                            )
-                        with gr.Row():
-                            top_p = gr.Slider(
-                                minimum=0.1, maximum=1.0, value=self.cfg.model.top_p,
-                                step=0.05, label="Top-p (تنوع)"
-                            )
-                            num_beams = gr.Slider(
-                                minimum=1, maximum=8, value=self.cfg.model.num_beams,
-                                step=1, label="تعداد Beam"
-                            )
-
-                    # Input/Output
-                    with gr.Row():
-                        with gr.Column(scale=1):
-                            question_input = gr.Textbox(
-                                label="سوال حقوقی خود را وارد کنید",
-                                placeholder="مثال: شرایط فسخ قرارداد اجاره چیست؟",
-                                lines=3
-                            )
-                            submit_btn = gr.Button("🔍 دریافت پاسخ", variant="primary")
-                        with gr.Column(scale=1):
-                            response_output = gr.Textbox(
-                                label="پاسخ سیستم",
-                                lines=8,
-                                interactive=False
-                            )
-                            references_output = gr.Textbox(
-                                label="مراجع حقوقی مرتبط",
-                                lines=6,
-                                interactive=False
-                            )
-                            metrics_output = gr.Textbox(
-                                label="معیارهای عملکرد",
-                                lines=1,
-                                interactive=False,
-                                elem_classes=["metrics-box"]
-                            )
-
-                # Training Tab
-                with gr.Tab("🎓 آموزش مدل") as training_tab:
-                    with gr.Row():
-                        with gr.Column(scale=1):
-                            train_model_dropdown = gr.Dropdown(
-                                choices=model_choices,
-                                value=model_choices[0],
-                                label="انتخاب مدل برای آموزش"
-                            )
-                            use_rag_training_checkbox = gr.Checkbox(
-                                value=True,
-                                label="استفاده از RAG در آموزش",
-                                info="استفاده از مواد قانونی در آموزش"
-                            )
-                            train_file_upload = gr.File(
-                                label="بارگذاری فایل‌ها�� آموزشی (JSONL)",
-                                file_types=[".jsonl"],
-                                type="filepath",
-                                file_count="multiple"
-                            )
-                        with gr.Column(scale=1):
-                            with gr.Accordion("⚙️ پارامترهای آموزش", open=False):
-                                train_epochs = gr.Slider(
-                                    minimum=1, maximum=10, value=self.cfg.epochs,
-                                    step=1, label="تعداد Epoch"
-                                )
-                                train_batch_size = gr.Slider(
-                                    minimum=1, maximum=16, value=self.cfg.batch_size,
-                                    step=1, label="اندازه Batch"
-                                )
-                                train_lr = gr.Slider(
-                                    minimum=1e-6, maximum=1e-3, value=self.cfg.lr,
-                                    step=1e-5, label="نرخ یادگیری"
-                                )
-
-                    train_btn = gr.Button("🎯 شروع آموزش", variant="primary")
-                    train_status = gr.Textbox(
-                        label="وضعیت آموزش",
-                        interactive=False,
-                        elem_classes=["status-box"]
-                    )
-                    train_progress = gr.Progress(label="پیشرفت آموزش")
-
-            # Event handlers
-            load_model_btn.click(
-                fn=lambda m, r: self.handle_load_model(m, r),
-                inputs=[model_dropdown, use_rag_checkbox],
-                outputs=load_status
-            )
-
-            submit_btn.click(
-                fn=lambda q, r, f, m, t, p, b: self.handle_generate_response(
-                    q, r, f, m, t, p, b
-                ),
-                inputs=[
-                    question_input,
-                    use_rag_checkbox,
-                    use_formalizer_checkbox,
-                    max_new_tokens,
-                    temperature,
-                    top_p,
-                    num_beams
-                ],
-                outputs=[response_output, references_output, metrics_output]
-            )
-
-            refresh_status_btn.click(
-                fn=lambda: self.get_system_status(),
-                outputs=system_status
-            )
+                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, 8, value=self.scfg.train.epochs, step=1, label="epochs")
+                    batch = gr.Slider(1, 16, value=self.scfg.train.batch_size, step=1, label="batch per device")
+                    lr = gr.Number(value=self.scfg.train.lr, label="learning rate")
+                with gr.Row():
+                    wd = gr.Number(value=self.scfg.train.weight_decay, label="weight decay")
+                    warmup = gr.Slider(0.0, 0.2, value=self.scfg.train.warmup_ratio, step=0.01, label="warmup ratio")
+                    report_to = gr.Dropdown(choices=["none","wandb"], value=self.scfg.train.report_to, label="report_to")
+                train_btn = gr.Button("شروع آموزش", variant="primary")
+                train_status = gr.Textbox(label="وضعیت آموزش", interactive=False)
+
+            # رویدادها
+            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])
+
+            index_btn.click(lambda f, ik, tk: self.build_index(f, ik, tk),
+                            inputs=[laws_file, id_key, text_key], outputs=index_status)
 
             train_btn.click(
-                fn=lambda m, f, r, e, b, lr, p, s: self.handle_training(
-                    m, f, r, e, b, lr, p, s
-                ),
-                inputs=[
-                    train_model_dropdown,
-                    train_file_upload,
-                    use_rag_training_checkbox,
-                    train_epochs,
-                    train_batch_size,
-                    train_lr,
-                    train_progress,
-                    train_status
-                ],
+                lambda choice, files, rag, e, b, l, _wd, _wu, _r:
+                    self.train(*_resolve(choice), files, rag, e, b, l, _wd, _wu, _r),
+                inputs=[model_dd_train, train_files, use_rag_train, epochs, batch, lr, wd, warmup, report_to],
                 outputs=train_status
             )
-
         return app
 
 # ==========================
-# Main Application
+# Entrypoint for HF Spaces
 # ==========================
-def main():
-    """Main entry point for the application"""
-    # Initialize system
+if __name__ == "__main__":
     app = LegalApp()
-
-    # Build and launch UI
     ui = app.build_ui()
-    ui.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        inbrowser=True,
-        share=False
-    )
-
-if __name__ == "__main__":
-    main()
+    ui.queue(concurrency_count=2).launch(server_name="0.0.0.0", server_port=7860)