Upload app(8).py

app(8).py

@@ -0,0 +1,1616 @@
+# -*- 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
+"""
+
+from __future__ import annotations
+import os
+import json
+import warnings
+import hashlib
+import threading
+import time
+from contextlib import contextmanager
+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
+
+import torch
+from torch.utils.data import Dataset
+from sklearn.model_selection import train_test_split
+from pydantic import BaseModel, validator, Field
+
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM,
+    AutoModelForCausalLM,
+    Trainer,
+    TrainingArguments,
+    EarlyStoppingCallback,
+    DataCollatorForSeq2Seq,
+    TrainerCallback
+)
+
+import chromadb
+from sentence_transformers import SentenceTransformer
+import gradio as gr
+
+warnings.filterwarnings("ignore")
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# ==========================
+# Enhanced Config with Validation
+# ==========================
+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
+
+class SystemConfig(BaseModel):
+    model: ModelConfig = Field(default_factory=ModelConfig)
+    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"
+    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
+
+# ==========================
+# 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()
+
+# ==========================
+# Enhanced Utilities
+# ==========================
+def set_seed_all(seed: int = 42):
+    import random
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    logger.info(f"Set random seed to {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"""
+    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} خط)"
+
+    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")
+
+# ==========================
+# Enhanced RAG System
+# ==========================
+class LegalRAGSystem:
+    def __init__(self, cfg: SystemConfig):
+        self.cfg = cfg
+        self.embedding_model: Optional[SentenceTransformer] = None
+        self.client = None
+        self.collection = None
+        self._lock = threading.Lock()
+
+    @contextmanager
+    def _safe_operation(self, operation_name: str):
+        """Context manager for safe RAG operations"""
+        start_time = time.time()
+        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"):
+            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)
+                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)}"
+
+    def retrieve(self, query: str) -> List[Dict]:
+        if not self.collection or not query.strip():
+            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
+
+        except Exception as e:
+            logger.error(f"Text formalization failed: {e}")
+            return text  # Return original text on error
+
+# ==========================
+# Enhanced Model Loader
+# ==========================
+class ModelLoader:
+    def __init__(self, model_config: ModelConfig):
+        self.cfg = model_config
+        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
+        if torch.cuda.is_available():
+            initial_memory = torch.cuda.memory_allocated()
+
+        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():
+            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 = 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,
+                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,
+                max_new_tokens=self.cfg.max_new_tokens,
+                do_sample=True,
+                temperature=max(0.1, self.cfg.temperature),  # Ensure min 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
+            )
+
+        # 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
+
+# ==========================
+# Enhanced 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
+        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")
+
+    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()
+        }
+
+class CausalJSONLDataset(Dataset):
+    def __init__(self, data: List[Dict], tokenizer, max_length: int):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+        # Process data
+        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")
+
+    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()
+
+        labels = input_ids.clone()
+        labels[attention_mask == 0] = -100
+
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels
+        }
+
+# ==========================
+# Enhanced Progress Callback
+# ==========================
+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}"
+        )
+
+# ==========================
+# Enhanced 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)
+
+            return True, "مدل Causal با موفقیت آموزش داده شد"
+
+        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
+            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
+        )
+
+    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}")
+
+# ==========================
+# Enhanced Legal App
+# ==========================
+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:
+            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})"
+
+            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 "مدل نامعتبر انتخاب شده"
+
+            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()
+
+        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
+
+        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)
+
+        except Exception as e:
+            return f"خطا در دریافت وضعیت: {str(e)}"
+
+    def _get_model_configs(self) -> Dict[str, Tuple[str, str]]:
+        """Get available model configurations"""
+        return {
+            "Seq2Seq (parsi-t5-base)": ("persiannlp/parsi-t5-base", "seq2seq"),
+            "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"),
+        }
+
+    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:
+
+            gr.HTML("""
+            <div style='text-align: center; margin-bottom: 20px;'>
+                <h1>ماحون — مشاور حقوقی هوشمند 🏛️</h1>
+                <p>سیستم پیشرفته مشاوره حقوقی با قابلیت RAG، Fine-tuning و هوش مصنوعی</p>
+            </div>
+            """)
+
+            # System Status
+            with gr.Accordion("وضعیت سیستم", open=False):
+                system_status = gr.Markdown(
+                    value=self.get_system_status(),
+                    elem_classes=["status-box"]
+                )
+                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
+            )
+
+            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
+                ],
+                outputs=train_status
+            )
+
+        return app
+
+# ==========================
+# Main Application
+# ==========================
+def main():
+    """Main entry point for the application"""
+    # Initialize system
+    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()