feat: Add AraSpell spelling correction endpoint (POST /api/spelling) with lazy model loading - Extract AraSpell pipeline into src/nlp/spelling/ (araspell_rules.py + araspell_service.py) - Pre-download AraBERT encoder-decoder + checkpoint in Dockerfile - Lazy load on first request, contextual=False to save RAM - Update health check to report spelling status

src/ara_spell.py → AraSpell.py

@@ -1,35 +1,54 @@
-# ## 📦 Part 1: Imports & Setup
-
+# AraSpell — Arabic Spell Checker Pipeline
+# Production-ready version
 
 import re
+import math
+import logging
 import torch
 import os
+from collections import Counter
 from transformers import AutoTokenizer, EncoderDecoderModel
 import Levenshtein
-import jellyfish  # NEW: For Damerau-Levenshtein (transpositions as 1 edit)
-
-print("✅ All imports successful")
-print(f"🔧 PyTorch version: {torch.__version__}")
-print(f"🔧 CUDA available: {torch.cuda.is_available()}")
-if torch.cuda.is_available():
-    print(f"🔧 GPU: {torch.cuda.get_device_name(0)}")
+import jellyfish
 
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
 
 # ═══════════════════════════════════════════════════════════════════════════════
-# 🏆 PRODUCTION SPELL CHECKER - OPTIMIZED VERSION (Merged from src/)
+# LOAD ARABERT SEQ2SEQ MODEL
 # ═══════════════════════════════════════════════════════════════════════════════
 
+from huggingface_hub import hf_hub_download
 
-# ═══════════════════════════════════════════════════════════════════════════════
-# 🏆 PRODUCTION SPELL CHECKER - OPTIMIZED VERSION (Merged from src/)
-# ═══════════════════════════════════════════════════════════════════════════════
-# IMPROVEMENTS:
-# 1. unified_collapse_repeated() - More conservative (3+ for Arabic)
-# 2. fix_hamza_conservative() - Only fixes word endings
-# 3. remove_hallucinations() - Removes duplicate words and trailing 'و'
-# 4. Expanded PREPOSITIONS - 16 prepositions instead of 2
-# 5. Better word splitting and joining
-# ═══════════════════════════════════════════════════════════════════════════════
+MODEL_REPO = 'bayan10/AraSpell-Model'
+MODEL_FILENAME = 'last_model.pt'
+
+try:
+    logger.info(f"Downloading/loading model from Hugging Face: {MODEL_REPO}")
+    MODEL_PATH = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
+except Exception as e:
+    raise RuntimeError(f"Failed to download model from Hugging Face: {e}")
+
+MODEL_NAME = 'aubmindlab/bert-base-arabertv02'
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = EncoderDecoderModel.from_encoder_decoder_pretrained(MODEL_NAME, MODEL_NAME)
+
+model.config.decoder_start_token_id = tokenizer.cls_token_id
+model.config.pad_token_id = tokenizer.pad_token_id
+model.config.eos_token_id = tokenizer.sep_token_id
+model.generation_config.max_length = 128
+model.generation_config.decoder_start_token_id = tokenizer.cls_token_id
+model.generation_config.pad_token_id = tokenizer.pad_token_id
+model.generation_config.eos_token_id = tokenizer.sep_token_id
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+checkpoint = torch.load(MODEL_PATH, map_location=device, weights_only=False)
+model.load_state_dict(checkpoint['model_state_dict'], strict=False)
+model = model.to(device)
+model.eval()
+
+logger.info(f"Model loaded on {device}, epoch: {checkpoint.get('epoch', 'N/A')}")
 
 from enum import Enum
 from typing import List, Tuple, Optional
@@ -47,11 +66,11 @@ class ErrorType(Enum):
     CLEAN = "clean"
 
 # ═══════════════════════════════════════════════════════════════════════════════
-# POST PROCESSOR (OPTIMIZED - Merged from src/)
+# POST PROCESSOR
 # ═══════════════════════════════════════════════════════════════════════════════
 
 class AraSpellPostProcessor:
-    """Post-processing techniques (OPTIMIZED VERSION)"""
+    """Arabic text post-processing techniques."""
     
     ARABIC_HARAKAT = 'ًٌٍَُِّْ'
     TATWEEL = 'ـ'
@@ -60,9 +79,7 @@ class AraSpellPostProcessor:
     }
     ARABIC_CONSONANTS = set('بتثجحخدذرزسشصضطظعغفقكلمن')
     
-    # ═══════════════════════════════════════════════════════════════════════════════
-    # BASIC NORMALIZATION
-    # ═══════════════════════════════════════════════════════════════════════════════
+    # --- Basic Normalization ---
     
     @staticmethod
     def remove_harakat(text: str) -> str:
@@ -81,16 +98,13 @@ class AraSpellPostProcessor:
             text = text.replace(old, new)
         return text
     
-    # ═══════════════════════════════════════════════════════════════════════════════
-    # UNIFIED CORE FUNCTIONS (NEW from src/)
-    # ═══════════════════════════════════════════════════════════════════════════════
+    # --- Core Functions ---
     
     @staticmethod
     def unified_collapse_repeated(text: str) -> str:
         """
-        UNIFIED repetition collapse (IMPROVED!)
-        - Arabic: 3+ consecutive → 1 (more conservative!)
-        - Latin: 2+ consecutive → 1
+        Collapse repeated characters.
+        Arabic: 3+ consecutive → 1 | Latin: 2+ consecutive → 1
         """
         # Arabic characters: 3+ → 1
         text = re.sub(r"([\u0600-\u06FF])\1{2,}", r"\1", text)
@@ -100,14 +114,9 @@ class AraSpellPostProcessor:
         
         return text
     
-
-    
     @staticmethod
     def remove_duplicate_words(text: str) -> str:
-        """
-        Remove consecutive duplicate words (NEW from src/)
-        Examples: كتاب كتاب → كتاب
-        """
+        """Remove consecutive duplicate words. e.g. كتاب كتاب → كتاب"""
         words = text.split()
         if len(words) < 2:
             return text
@@ -119,13 +128,9 @@ class AraSpellPostProcessor:
         
         return ' '.join(result)
     
-
-    
     @staticmethod
     def normalize_spaces(text: str) -> str:
-        """
-        UNIFIED space normalization (NEW from src/)
-        """
+        """Normalize whitespace: multiple spaces, unicode spaces, punctuation spacing."""
         # Multiple spaces → single
         text = re.sub(r' +', ' ', text)
         
@@ -159,20 +164,11 @@ class AraSpellPostProcessor:
                 i += 1
         return ' '.join(result)
     
-    # ═══════════════════════════════════════════════════════════════════════════════
-    # HAMZA HANDLING (NEW from src/)
-    # ═══════════════════════════════════════════════════════════════════════════════
+    # --- Hamza & Ta Marbuta Handling ---
     
     @staticmethod
     def fix_hamza_conservative(text: str) -> str:
-        """
-        CONSERVATIVE Hamza normalization (NEW!)
-        Only normalizes at word END, not middle
-        
-        Examples:
-        ✓ "المدرسه" → "المدرسة" (end of word)
-        ✓ "سأل" → "سأل" (middle - KEEP IT!)
-        """
+        """Conservative Hamza normalization — only at word END, not middle."""
         words = text.split()
         result = []
         
@@ -191,34 +187,63 @@ class AraSpellPostProcessor:
         return ' '.join(result)
     
     @staticmethod
-    def fix_ha_ta_marbuta(text: str) -> str:
+    def fix_ha_ta_marbuta(text: str, vocab_manager=None) -> str:
         """
-        Fix ه → ة at end of words (pattern-based)
+        Smart ه → ة fix at end of words.
+        
+        Key insight: ه at word end can be:
+        - Ta Marbuta (should be ة): المدرسه → المدرسة
+        - Possessive pronoun (should stay ه): تحقيقه = his achievement
+        
+        Strategy: Only convert if the ة version is IV (in tokenizer vocab).
+        This distinguishes المدرسة (IV) from تحقيقة (not a real word form).
+        Without vocab_manager, falls back to original pattern-based approach.
         """
+        # Protected words: anything containing لله
+        PROTECTED_ENDINGS = ['لله']
+        
         words = text.split()
         result = []
         
         for word in words:
+            # Skip protected words (Allah-related)
+            if any(word.endswith(e) for e in PROTECTED_ENDINGS):
+                result.append(word)
+                continue
+            
             if len(word) >= 4 and word.endswith('ه'):
                 # Check if second-to-last char is a consonant
                 if word[-2] in AraSpellPostProcessor.ARABIC_CONSONANTS:
-                    result.append(word[:-1] + 'ة')
-                    continue
+                    candidate_with_ta = word[:-1] + 'ة'
+                    
+                    if vocab_manager:
+                        # SMART MODE: Use vocab to decide
+                        ta_iv = vocab_manager.is_iv(candidate_with_ta)
+                        ha_iv = vocab_manager.is_iv(word)
+                        
+                        if ta_iv:
+                            # ة version is IV → convert (المدرسه→المدرسة)
+                            result.append(candidate_with_ta)
+                            continue
+                        elif ha_iv:
+                            # Only ه version is IV → keep ه (possessive: تحقيقه)
+                            result.append(word)
+                            continue
+                        # else: NEITHER is IV → keep original ه
+                        # (safer than guessing — could be rare possessive)
+                    else:
+                        # FALLBACK: No vocab → use original pattern-based approach
+                        result.append(candidate_with_ta)
+                        continue
             result.append(word)
         
         return ' '.join(result)
     
-    # ═══════════════════════════════════════════════════════════════════════════════
-    # HALLUCINATION REMOVAL (NEW from src/)
-    # ═══════════════════════════════════════════════════════════════════════════════
+    # --- Hallucination Removal ---
     
     @staticmethod
     def remove_hallucinations(text: str) -> str:
-        """
-        Remove model hallucinations (NEW from src/):
-        - Duplicate words
-        - Trailing 'و' artifacts
-        """
+        """Remove model hallucinations: duplicate words, trailing 'و' artifacts."""
         words = text.split()
         if not words:
             return text
@@ -270,9 +295,7 @@ class AraSpellPostProcessor:
         
         return text
     
-    # ═══════════════════════════════════════════════════════════════════════════════
-    # WORD SPLITTING & MERGING (NEW from src/)
-    # ═══════════════════════════════════════════════════════════════════════════════
+    # --- Word Splitting & Merging ---
     
     @staticmethod
     def merge_separated_al(text: str) -> str:
@@ -281,12 +304,7 @@ class AraSpellPostProcessor:
     
     @staticmethod
     def join_fragments(text: str) -> str:
-        """
-        IMPROVED: Join short fragments with better validation (NEW from src/)
-        الط + الب → الطالب
-        
-        FIXED: No longer joins valid separate words like "في من"
-        """
+        """Join short fragments with validation. e.g. الط + الب → الطالب"""
         words = text.split()
         if len(words) < 2:
             return text
@@ -341,14 +359,13 @@ class AraSpellPostProcessor:
         
         return ' '.join(result)
     
-    # ═══════════════════════════════════════════════════════════════════════════════
-    # MAIN PIPELINES
-    # ═══════════════════════════════════════════════════════════════════════════════
+    # --- Main Pipelines ---
     
     @staticmethod
-    def full_postprocess(text: str, original: str = "") -> str:
+    def full_postprocess(text: str, original: str = "", vocab_manager=None) -> str:
         """
         Apply all post-processing steps (OPTIMIZED ORDER!)
+        vocab_manager: optional, enables smart ه/ة handling
         """
         # 1. Remove hallucinated prefixes
         if original:
@@ -366,8 +383,8 @@ class AraSpellPostProcessor:
         # 5. Fix Hamza (CONSERVATIVE!)
         text = AraSpellPostProcessor.fix_hamza_conservative(text)
         
-        # 6. Fix Ta Marbuta
-        text = AraSpellPostProcessor.fix_ha_ta_marbuta(text)
+        # 6. Fix Ta Marbuta (SMART MODE with vocab_manager!)
+        text = AraSpellPostProcessor.fix_ha_ta_marbuta(text, vocab_manager=vocab_manager)
         
         # 7. Remove word repetition with 'و'
         text = AraSpellPostProcessor.remove_word_repetition_with_wa(text)
@@ -382,7 +399,7 @@ class AraSpellPostProcessor:
 
 
 # ─────────────────────────────────────────────────────────────────────────────
-# ERROR CLASSIFIER (Keep from original notebook)
+# ERROR CLASSIFIER
 # ─────────────────────────────────────────────────────────────────────────────
 
 class ErrorClassifier:
@@ -427,18 +444,12 @@ class ErrorClassifier:
         else:
             return ErrorType.CLEAN
 
-
 # ═══════════════════════════════════════════════════════════════════════════════
-# Part 3: Production Spell Checker Class
-# This is the best pipeline based on extensive testing of 8+ different approaches
-# ═══════════════════════════════════════════════════════════════════════════════
-
-# ═══════════════════════════════════════════════════════════════════════════════
-# RULES-BASED CORRECTOR (EXPANDED - Merged from src/)
+# RULES-BASED CORRECTOR
 # ═══════════════════════════════════════════════════════════════════════════════
 
 class RulesBasedCorrector:
-    """Rules-based correction (EXPANDED VERSION)"""
+    """Rules-based correction with keyboard proximity mapping."""
     
     # Persian/Urdu → Arabic mapping
     SUBSTITUTION_MAP = {
@@ -459,6 +470,55 @@ class RulesBasedCorrector:
         'لل'
     }
     
+    # Keyboard Proximity Mapping
+    # Arabic keyboard layout adjacency
+    KEYBOARD_NEIGHBORS = {
+        'ض': ['ص', 'ق'],
+        'ص': ['ض', 'ث', 'ق'],
+        'ث': ['ص', 'ق'],
+        'ق': ['ض', 'ص', 'ث', 'ف', 'غ'],
+        'ف': ['ق', 'غ', 'ع', 'ب'],
+        'غ': ['ق', 'ف', 'ع', 'ه'],
+        'ع': ['ف', 'غ', 'ه', 'خ'],
+        'ه': ['غ', 'ع', 'خ', 'ح'],
+        'خ': ['ع', 'ه', 'ح', 'ج'],
+        'ح': ['ه', 'خ', 'ج'],
+        'ج': ['خ', 'ح', 'د'],
+        'د': ['ج', 'ذ'],
+        'ذ': ['د'],
+        'ش': ['س', 'ي', 'ئ'],
+        'س': ['ش', 'ي', 'ب'],
+        'ي': ['ش', 'س', 'ب', 'ت'],
+        'ب': ['ي', 'س', 'ف', 'ل', 'ن'],
+        'ل': ['ب', 'ا', 'ن', 'م'],
+        'ا': ['ل', 'ت', 'م'],
+        'ت': ['ي', 'ا', 'ن'],
+        'ن': ['ب', 'ل', 'ت', 'م', 'ك'],
+        'م': ['ل', 'ا', 'ن', 'ك'],
+        'ك': ['ن', 'م', 'ط'],
+        'ط': ['ك', 'ظ'],
+        'ظ': ['ط'],
+        'ئ': ['ش', 'ء', 'ر'],
+        'ء': ['ئ', 'ؤ'],
+        'ؤ': ['ء', 'ر'],
+        'ر': ['ئ', 'ؤ', 'لا', 'ى', 'ز'],
+        'لا': ['ر', 'ى'],
+        'ى': ['ر', 'لا', 'ة', 'ز'],
+        'ة': ['ى', 'و', 'ز'],
+        'و': ['ة', 'ز'],
+        'ز': ['ر', 'ى', 'ة', 'و'],
+        # Alif variants
+        'أ': ['ا', 'إ', 'آ'],
+        'إ': ['ا', 'أ'],
+        'آ': ['ا', 'أ'],
+    }
+    
+    @staticmethod
+    def is_keyboard_neighbor(char1: str, char2: str) -> bool:
+        """Check if two Arabic chars are adjacent on keyboard."""
+        neighbors = RulesBasedCorrector.KEYBOARD_NEIGHBORS.get(char1, [])
+        return char2 in neighbors
+    
     @staticmethod
     def fix_char_substitution(text: str) -> str:
         """Replace Persian/Urdu characters with Arabic"""
@@ -468,10 +528,7 @@ class RulesBasedCorrector:
     
     @staticmethod
     def fix_char_repetition(text: str) -> str:
-        """
-        Remove excessive character repetition (IMPROVED!)
-        Now: 3+ consecutive → 1 (more conservative)
-        """
+        """Remove excessive character repetition (3+ consecutive → 1)."""
         # Only collapse 3+ repetitions (not 2+)
         text = re.sub(r'([^\d\s])\1{2,}', r'\1', text)
         return text
@@ -596,11 +653,7 @@ class OutputValidator:
         if not corrected or not corrected.strip():
             return False, "empty_output"
         
-        # ═══════════════════════════════════════════════════════════════════════════
-        # SPACE LENIENCY (Phase 1 - Solutions.md الفكرة 3)
-        # ═══════════════════════════════════════════════════════════════════════════
-        # If the ONLY difference is whitespace → accept if resulting words are valid
-        # This fixes split/merge issues like "في��لمدرسة" → "في المدرسة"
+        # Space Leniency: if ONLY difference is whitespace → accept
         original_no_space = original.replace(' ', '').replace('\u200c', '')  # Also handle ZWNJ
         corrected_no_space = corrected.replace(' ', '').replace('\u200c', '')
         
@@ -639,9 +692,8 @@ class OutputValidator:
              
         return True, "valid"
 
-
 # ═══════════════════════════════════════════════════════════════════════════════
-# VOCABULARY MANAGER (NEW - Phase 1)
+# VOCABULARY MANAGER
 # ═══════════════════════════════════════════════════════════════════════════════
 
 class VocabularyManager:
@@ -673,7 +725,7 @@ class VocabularyManager:
         # Build normalized vocabulary for fuzzy matching
         self.normalized_vocab = {self.normalize_for_comparison(w): w for w in self.vocab}
         
-        print(f"📚 VocabularyManager initialized: {len(self.vocab)} words")
+        logger.info(f"VocabularyManager initialized: {len(self.vocab)} words")
     
     @classmethod
     def normalize_for_comparison(cls, word: str) -> str:
@@ -763,9 +815,8 @@ class VocabularyManager:
         max_len = max(len(original), len(corrected), 1)
         return 1.0 - (dist / max_len)
 
-
 # ═══════════════════════════════════════════════════════════════════════════════
-# WORD ALIGNER (Phase 2 - Solutions.md الفكرة 5)
+# WORD ALIGNER
 # ═══════════════════════════════════════════════════════════════════════════════
 
 class WordAligner:
@@ -848,12 +899,26 @@ class WordAligner:
             return input_word 
             
         # Case 4: Both OOV
-        # Take output, usually closer to target even if still OOV
+        # Subword-level correction
+        # If words are similar length, try character-level blending to find IV
+        if len(input_word) == len(output_word) and len(input_word) >= 3:
+            # Try replacing one char at a time from output into input
+            for i in range(len(input_word)):
+                if input_word[i] != output_word[i]:
+                    # Try input with this one char from output
+                    hybrid = input_word[:i] + output_word[i] + input_word[i+1:]
+                    if self.vocab.is_iv(hybrid):
+                        return hybrid
+                    # Try output with this one char from input
+                    hybrid2 = output_word[:i] + input_word[i] + output_word[i+1:]
+                    if self.vocab.is_iv(hybrid2):
+                        return hybrid2
+        
+        # Default: Take output, usually closer to target even if still OOV
         return output_word
 
-
 # ═══════════════════════════════════════════════════════════════════════════════
-# SPLIT/MERGE SPECIALIST (Phase 2 - Solutions.md الفكرة 4)
+# SPLIT/MERGE SPECIALIST
 # ═══════════════════════════════════════════════════════════════════════════════
 
 class SplitMergeSpecialist:
@@ -893,15 +958,24 @@ class SplitMergeSpecialist:
             self.SEPARABLE_PREFIXES, key=len, reverse=True
         )
     
+    # Attached prefix patterns that should NOT be split (normal Arabic word formations)
+    ATTACHED_PREFIXES = [
+        'وال', 'بال', 'فال', 'كال', 'لل',   # Conjunction/Preposition + Article
+        'وب', 'وف', 'ول', 'وك', 'وم', 'ون',  # Conjunction + Preposition
+        'فب', 'فل', 'فك', 'فم',              # Conjunction + Preposition
+    ]
+    
     def split_word(self, word: str) -> str:
         """
         Try to split an OOV word into IV components.
         
-        STRICT Strategy: Only split when BOTH parts are IV.
-        This prevents over-splitting like معظم → مع ظم
+        Strict Strategy:
+        - Only split when BOTH parts are IV
+        - Protect attached prefix patterns (وال، بال، etc.)
+        - Minimum part lengths to prevent micro-splits
         """
-        # Short words: don't split
-        if len(word) < 4:
+        # Short words: don't split (increased from 4 to 5 for safety)
+        if len(word) < 5:
             return word
         
         # Already IV: no need to split
@@ -912,17 +986,29 @@ class SplitMergeSpecialist:
         if word in self.PROTECTED_WORDS:
             return word
         
+        # Protected prefix patterns (وال، بال، فال، etc.)
+        # These are normal Arabic word formations, NOT merge errors
+        for prefix in self.ATTACHED_PREFIXES:
+            if word.startswith(prefix):
+                remainder = word[len(prefix):]
+                # If the remainder (without the prefix) is IV, this is a valid prefixed word
+                if self.vocab.is_iv(remainder):
+                    return word  # Don't split — it's prefix+valid_word
+                # Also check with article: e.g. والخصوصي → وال+خصوصي, check خصوصي
+                if prefix.endswith('ال') and self.vocab.is_iv(remainder):
+                    return word
+        
         # 1. Try separable prefixes first (higher priority)
         for prefix in self.separable_prefixes:
-            if word.startswith(prefix) and len(word) > len(prefix) + 1:
+            if word.startswith(prefix) and len(word) > len(prefix) + 2:  # Remainder must be > 2 chars
                 remainder = word[len(prefix):]
                 
                 # Only accept if remainder is IV
                 if self.vocab.is_iv(remainder):
                     return f"{prefix} {remainder}"
         
-        # 2. Try all positions - STRICT: BOTH parts must be IV
-        for i in range(2, len(word) - 1):
+        # 2. Try all positions - STRICT: BOTH parts must be IV AND both >= 3 chars
+        for i in range(3, len(word) - 2):  # Both parts at least 3 chars
             left = word[:i]
             right = word[i:]
             
@@ -932,14 +1018,20 @@ class SplitMergeSpecialist:
         # No valid split found
         return word
     
+    # Common Arabic pronoun/possessive suffixes (2-3 chars)
+    # These are often incorrectly split from their host word
+    PRONOUN_SUFFIXES = {'كم', 'هم', 'ها', 'هن', 'كن', 'نا', 'هما', 'كما', 'تم', 'تن'}
+    
     def merge_fragments(self, text: str) -> str:
         """
         Try to merge adjacent OOV fragments into IV words.
         
         Key patterns:
-        1. Ta-marbuta detachment: السوري ة → السورية (Safe even if السوري is IV)
+        1. Ta-marbuta detachment: السوري ة → السورية
         2. Al- detachment: ال كتاب → الكتاب
         3. General OOV+OOV merging: Only if both are OOV and result is IV
+        4. Short OOV fragment: 1-2 char OOV + next → IV
+        5. Pronoun suffix reattachment: علي كم → عليكم
         """
         words = text.split()
         if len(words) < 2:
@@ -986,6 +1078,24 @@ class SplitMergeSpecialist:
                         result.append(merged)
                         i += 2
                         continue
+                
+                # Pattern 5: Pronoun suffix reattachment
+                # Fixes over-splitting: علي كم → عليكم
+                if next_word in self.PRONOUN_SUFFIXES:
+                    if self.vocab.is_iv(merged) and not self.vocab.is_iv(word):
+                        result.append(merged)
+                        i += 2
+                        continue
+                
+                # Pattern 6: Short fragment merge
+                # Merges two short words when combined they form a valid longer word
+                # Fixes: علي كم → عليكم, ويت أمل → ويتأمل, المد فتر → المدفتر
+                # Condition: both words ≤ 3 chars, merged ≥ 5 chars and IV
+                if len(word) <= 3 and len(next_word) <= 3:
+                    if len(merged) >= 5 and self.vocab.is_iv(merged):
+                        result.append(merged)
+                        i += 2
+                        continue
             
             result.append(word)
             i += 1
@@ -1013,12 +1123,10 @@ class SplitMergeSpecialist:
         
         return ' '.join(processed)
 
-
 # ═══════════════════════════════════════════════════════════════════════════════
-# EDIT DISTANCE CORRECTOR (NEW!)
+# EDIT DISTANCE CORRECTOR
 # ═══════════════════════════════════════════════════════════════════════════════
 
-
 class EditDistanceCorrector:
     """
     Generates candidates based on Levenshtein distance.
@@ -1089,154 +1197,7 @@ class EditDistanceCorrector:
         return ' '.join(corrected_words)
 
     
-    @staticmethod
-    def check_word_count(original: str, corrected: str) -> Tuple[bool, str]:
-        """Check if word count is reasonable"""
-        words_original = original.split()
-        words_corrected = corrected.split()
-        
-        # Allow more flexibility for merged words
-        # If original has long words (>7 chars), allow more splits
-        has_long_word = any(len(w) > 7 for w in words_original)
-        max_expansion = 3 if has_long_word else 1
-        
-        if len(words_corrected) > len(words_original) + max_expansion:
-            return False, "too_many_words"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_token_overlap(original: str, corrected: str) -> Tuple[bool, str]:
-        """Check token overlap"""
-        tokens_original = set(original.split())
-        tokens_corrected = set(corrected.split())
-        
-        if not tokens_original:
-            return True, "valid"
-            
-        # Skip for short sentences (1-2 words) where edit distance is better
-        if len(tokens_original) <= 2:
-            return True, "valid"
-        
-        overlap = len(tokens_original & tokens_corrected) / len(tokens_original)
-        
-        if overlap < 0.2:
-            return False, "low_token_overlap"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_edit_distance_reasonable(original: str, corrected: str) -> Tuple[bool, str]:
-        """Check if edit distance is reasonable"""
-        distance = OutputValidator.calculate_edit_distance(original, corrected)
-        threshold = len(original) * 0.5  # Adaptive threshold
-        
-        if distance > threshold:
-            return False, "excessive_edit_distance"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_length(original: str, corrected: str) -> Tuple[bool, str]:
-        """Check if length change is reasonable"""
-        max_change = len(original) * 0.25 + 3
-        
-        if abs(len(corrected) - len(original)) > max_change:
-            return False, "excessive_length_change"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_repetition(corrected: str) -> Tuple[bool, str]:
-        """Check for suspicious repetitions"""
-        # Check for 4+ consecutive identical chars
-        if re.search(r'(.)\1{3,}', corrected):
-            return False, "excessive_repetition"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_word_quality(corrected: str) -> Tuple[bool, str]:
-        """Check for suspicious short words"""
-        words = corrected.split()
-        
-        # Count 1-char words
-        single_char_count = sum(1 for w in words if len(w) == 1)
-        
-        if single_char_count > 2:
-            return False, "too_many_short_words"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_word_splitting_quality(corrected: str) -> Tuple[bool, str]:
-        """Check if word splitting looks suspicious"""
-        words = corrected.split()
-        
-        # Check for many 2-char words
-        two_char_count = sum(1 for w in words if len(w) == 2)
-        
-        if len(words) > 3 and two_char_count > len(words) * 0.5:
-            return False, "suspicious_word_splitting"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_single_chars(corrected: str) -> Tuple[bool, str]:
-        """Check for suspicious single characters"""
-        words = corrected.split()
-        
-        for word in words:
-            if len(word) == 1 and word not in {'و', 'ب', 'ل', 'ك', 'ف'}:
-                return False, "suspicious_short_word"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def check_word_preservation(original: str, corrected: str) -> Tuple[bool, str]:
-        """Check if at least some words are preserved"""
-        words_original = set(original.split())
-        words_corrected = set(corrected.split())
-        
-        if len(words_original) <= 1:
-            return True, "valid"
-            
-        # Skip if splitting a single merged word
-        if len(words_original) == 1 and len(words_corrected) > 1:
-            return True, "valid"
-        
-        preserved = len(words_original & words_corrected) / len(words_original)
-        
-        if preserved < 0.1:
-            return False, "too_few_preserved_words"
-        
-        return True, "valid"
-    
-    @staticmethod
-    def validate(original: str, corrected: str, error_type: str = "unknown") -> Tuple[bool, str]:
-        """
-        Validate model output
-        Returns: (is_valid, reason)
-        """
-        # Run all checks
-        checks = [
-            OutputValidator.check_character_preservation(original, corrected),
-            OutputValidator.check_word_count(original, corrected),
-            OutputValidator.check_token_overlap(original, corrected),
-            OutputValidator.check_edit_distance_reasonable(original, corrected),
-            OutputValidator.check_length(original, corrected),
-            OutputValidator.check_repetition(corrected),
-            OutputValidator.check_word_quality(corrected),
-            OutputValidator.check_word_splitting_quality(corrected),
-            OutputValidator.check_single_chars(corrected),
-            OutputValidator.check_word_preservation(original, corrected),
-        ]
-        
-        for is_valid, reason in checks:
-            if not is_valid:
-                return False, reason
-        
-        return True, "valid"
+
 
 
 
@@ -1310,7 +1271,7 @@ class ContextualCorrector:
                 if candidate not in candidates:
                     candidates.append(candidate)
         
-        # 2. 🆕 Remove repeated characters (deletion)
+        # 2. Remove repeated characters (deletion)
         # Fixes: مدررسة -> مدرسة, جميلل -> جميل
         for i in range(len(word) - 1):
             if word[i] == word[i+1]:
@@ -1319,7 +1280,7 @@ class ContextualCorrector:
                 if candidate not in candidates:
                     candidates.append(candidate)
         
-        # 3. 🆕 Edit Distance 1 Candidates (Insertions, Substitutions, Transpositions)
+        # 3. Edit Distance 1 Candidates (Insertions, Substitutions, Transpositions)
         # Using a restricted set of characters to avoid explosion
         COMMON_CHARS = 'ابتثجحخدذرزسشصضطظعغفقكلمنهويأإآءئؤةى'
         
@@ -1421,84 +1382,8 @@ class ContextualCorrector:
         
         return scores
     
-    def correct_word_in_context(self, text: str, position: int, threshold: float = 0.05) -> Tuple[str, dict]:
-        """
-        Correct a word in context
-        Returns: (corrected_word, metadata)
-        """
-        words = text.split()
-        if position >= len(words):
-            return words[position] if position < len(words) else "", {}
-        
-        original_word = words[position]
-        
-        # Generate candidates
-        candidates = self.generate_candidates(original_word)
-        
-        if len(candidates) == 1:
-            return original_word, {'candidates': 1, 'corrected': False}
-        
-        # Score candidates (batch)
-        scores = self.score_candidates_batch(text, position, candidates)
-        
-        # Find best candidate
-        best_word = max(scores, key=scores.get)
-        best_score = scores[best_word]
-        original_score = scores[original_word]
-        
-        # Check if original is in vocabulary (Single Token)
-        # We need to know if original_score is reliable (IV) or just a prefix score (OOV)
-        orig_tokens = self.tokenizer.encode(original_word, add_special_tokens=False)
-        is_orig_iv = len(orig_tokens) == 1
-        
-        # Apply correction logic
-        min_abs_improvement = 1e-4
-        is_improvement = False
-        
-        if is_orig_iv:
-            # Standard relative improvement for IV words
-            if best_score > original_score * (1 + threshold) and \
-               best_score > original_score + min_abs_improvement:
-                is_improvement = True
-        else:
-            # Original is OOV/Multi-token (likely specific scoring issue with prefixes like 'ال')
-            # If best_word is IV (Single Token), effectively compare its score against an absolute threshold
-            # because original_score (prefix) is misleadingly high.
-            # We use a stricter absolute threshold for this case to avoid over-correction of valid OOVs.
-            oov_correction_threshold = 0.001 # 0.1% probability minimum (Lowered to catch 'الطقس')
-            
-            # Also ensure best_word is IV (which we forced in generation, but good to be sure)
-            best_tokens = self.tokenizer.encode(best_word, add_special_tokens=False)
-            is_best_iv = len(best_tokens) == 1
-            
-            if is_best_iv and best_score > oov_correction_threshold:
-                # We ignore original_score here as it's likely just P(prefix)
-                is_improvement = True
-            elif best_score > original_score * (1 + threshold):
-                # Fallback to relative check if both are OOV or logic allows
-                is_improvement = True
-
-        if best_word != original_word and is_improvement:
-            return best_word, {
-                'candidates': len(candidates),
-                'corrected': True,
-                'original_score': original_score,
-                'best_score': best_score,
-                'improvement': best_score - original_score,
-                'is_orig_iv': is_orig_iv
-            }
-        
-        return original_word, {
-            'candidates': len(candidates),
-            'corrected': False,
-            'original_score': original_score
-        }
-    
     def predict_masked_token(self, text: str, position: int, top_k: int = 5) -> List[Tuple[str, float]]:
-        """
-        Predict words for a masked position.
-        Returns: List of (word, score)
-        """
+        """Predict words for a masked position. Returns list of (word, score)."""
         words = text.split()
         if position >= len(words):
             return []
@@ -1528,85 +1413,83 @@ class ContextualCorrector:
             token_id = top_k_indices[i].item()
             score = top_k_weights[i].item()
             token = self.tokenizer.decode([token_id]).strip()
-            # Filter out subwords (starting with ##) and special tokens
             if not token.startswith("##") and token not in self.tokenizer.all_special_tokens:
                 results.append((token, score))
                 
         return results
 
-    def refine_sentence_with_mask(self, text: str, threshold: float = 0.001) -> str:
-        """
-        Refine sentence by masking weak words and predicting replacements.
-        Effectively uses BERT as a contextual dictionary.
+    def refine_sentence_with_mask(self, text: str, threshold: float = 0.001, vocab_manager=None, raw_model_output=None) -> str:
+        """Refine sentence by masking weak words and predicting replacements.
+        IV-Safe + Strict similarity + BERT Kill Switch.
         """
         words = text.split()
         refined_words = words.copy()
         
+        # Build set of raw model words for kill switch
+        raw_words = raw_model_output.split() if raw_model_output else []
+        
         for i, word in enumerate(words):
+            # IV-Safe check - NEVER replace IV words
+            if vocab_manager and vocab_manager.is_iv(word):
+                continue
+            
+            # BERT Kill Switch: skip words matching raw model output
+            if i < len(raw_words) and word == raw_words[i]:
+                continue
+            
+            # Skip very short words (prepositions etc)
+            if len(word) <= 2:
+                continue
+            
             # 1. Check confidence
-            # We use score_with_mlm. If score is very low, it's a candidate for refinement.
             current_score = self.score_with_mlm(text, i, word)
             
-            # If word is confident enough, skip
-            # Threshold needs to be tuned. 0.001 implies 0.1% probability.
             if current_score > threshold:
                 continue
                 
             # 2. Mask and Predict
             predictions = self.predict_masked_token(text, i, top_k=10)
             
-            # 3. Filter and Select
+            # 3. Filter and Select (strict)
             for pred_word, pred_score in predictions:
                 if pred_word == word:
                     continue
 
-                # STRICTER CONSTRAINTS
-                
-                # 1. Length Check
-                if abs(len(pred_word) - len(word)) > 2:
+                if abs(len(pred_word) - len(word)) > 1:
                      continue
                      
-                # 2. Similarity Check (Edit Distance)
-                # We want to fix TYPOS, not semantic hallucinations.
-                # So the replacement MUST be structurally similar.
-                # Exception: If original word is very short (<3 chars), strict check.
-                
+                # Similarity Check (0.90 minimum)
                 dist = Levenshtein.distance(word, pred_word)
                 max_len = max(len(word), len(pred_word))
                 similarity = 1.0 - (dist / max_len)
                 
+                if similarity < 0.90:
+                    continue
                 
-                # Minimum similarity required: 0.7 (Much stricter to prevent semantic shift)
-                # 'بسرعّة' -> 'بسرعة' (High sim)
-                # 'الإسلامية' -> 'التطبيقية' (Low sim - rejected)
-                if similarity < 0.7:
+                # Must be IV
+                if vocab_manager and vocab_manager.is_oov(pred_word):
                     continue
-                    
-                # 3. Score Improvement
-                # If current word is total garbage (score < 1e-5), take any reasonable predicted word.
-                # If current word has some confidence, only replace if prediction is MUCH better.
                 
-                # Check if original word is IV and common
+                # Minimum absolute confidence gate (12%)
+                if pred_score < 0.12:
+                    continue
+                    
+                # Score Improvement
                 is_original_common = current_score > 0.001
                 
                 if is_original_common:
-                     # Very strict if original seems okay
-                     if pred_score > current_score * 500:
+                     if pred_score > current_score * 1000:
                          refined_words[i] = pred_word
                          break
                 else:
-                    # Looser if original is weak
-                    if pred_score > current_score * 10 or pred_score > 0.1:
+                    if pred_score > current_score * 50 and pred_score > 0.2:
                         refined_words[i] = pred_word
-                        break # Take the top valid prediction
+                        break
         
         return ' '.join(refined_words)
     
     def calculate_sentence_score(self, text: str) -> float:
-        """
-        Calculate a 'fluency' score for the sentence using BERT MLM.
-        Returns the average probability of each word being predicted in its context.
-        """
+        """Calculate fluency score using BERT MLM average word probability."""
         words = text.split()
         if not words:
             return 0.0
@@ -1615,8 +1498,6 @@ class ContextualCorrector:
         scored_words = 0
         
         for i, word in enumerate(words):
-            # We skip scoring very common stopwords to focus on content/structure?
-            # No, keep it simple for now: score everything.
             score = self.score_with_mlm(text, i, word)
             total_score += score
             scored_words += 1
@@ -1626,36 +1507,6 @@ class ContextualCorrector:
             
         return total_score / scored_words
 
-    def correct_sentence(self, text: str, threshold: float = 0.01) -> Tuple[str, dict]:
-        """
-        Correct all words in a sentence
-        Returns: (corrected_text, metadata)
-        """
-        words = text.split()
-        corrected_words = []
-        corrections = []
-        
-        for i, word in enumerate(words):
-            corrected_word, meta = self.correct_word_in_context(text, i, threshold)
-            corrected_words.append(corrected_word)
-            
-            if meta.get('corrected'):
-                corrections.append({
-                    'position': i,
-                    'original': word,
-                    'corrected': corrected_word,
-                    'confidence': meta.get('improvement', 0)
-                })
-        
-        corrected_text = ' '.join(corrected_words)
-        
-        return corrected_text, {
-            'corrections': corrections,
-            'total_words': len(words)
-        }
-
-
-
 
 # ═══════════════════════════════════════════════════════════════════════════════
 # MAIN SPELL CHECKER CLASS
@@ -1675,22 +1526,21 @@ class ArabicSpellChecker:
         self.classifier = ErrorClassifier()
         self.rules = RulesBasedCorrector()
         self.validator = OutputValidator()
-        self.vocab_manager = VocabularyManager(tokenizer)  # Phase 1: Vocabulary Manager
+        self.vocab_manager = VocabularyManager(tokenizer)
         self.edit_corrector = EditDistanceCorrector(tokenizer)  # Edit Distance candidates
-        self.split_merge = SplitMergeSpecialist(self.vocab_manager)  # Phase 2: Split/Merge
+        self.split_merge = SplitMergeSpecialist(self.vocab_manager)
         
-        # Phase 2: WordAligner for word-level hybrid corrections
+        # WordAligner for word-level hybrid corrections
         self.word_aligner = WordAligner(self.vocab_manager)
         
         # Initialize contextual corrector (optional)
         self.use_contextual = use_contextual
         if use_contextual:
             try:
-                # Initialize ContextualCorrector
                 self.contextual = ContextualCorrector()
-                print("✅ Contextual correction enabled")
+                logger.info("Contextual correction enabled")
             except Exception as e:
-                print(f"⚠️ Contextual correction disabled: {e}")
+                logger.warning(f"Contextual correction disabled: {e}")
                 self.contextual = None
                 self.use_contextual = False
         else:
@@ -1709,8 +1559,7 @@ class ArabicSpellChecker:
         return text
     
     def _fix_merged_with_errors(self, text: str) -> str:
-        """
-        🆕 Fix merged words that contain errors
+        """ Fix merged words that contain errors
         
         Examples:
             الممدرسة → المدرسة
@@ -1727,8 +1576,7 @@ class ArabicSpellChecker:
     
 
     def _split_merged_words_linguistic(self, text: str) -> str:
-        """
-        🆕 Split merged words using linguistic patterns
+        """ Split merged words using linguistic patterns
         
         Examples:
             كلصباح → كل صباح
@@ -1757,18 +1605,11 @@ class ArabicSpellChecker:
         
         # Pattern 6: على/عن in middle of (merged) words
         text = re.sub(r'([ا-ي]{3,})(على|عن)([ا-ي]{3,})', r'\1 \2 \3', text)
-
-        # Pattern 7: بسم الله الرحمن الرحيم (common concatenation)
-        text = re.sub(r'\bبسماللهالرحمنالرحيم\b', 'بسم الله الرحمن الرحيم', text)
-        text = re.sub(r'\bبسمالله\b', 'بسم الله', text)
-        text = re.sub(r'اللهالرحمن', 'الله الرحمن', text)
-        text = re.sub(r'الرحمنالرحيم', 'الرحمن الرحيم', text)
-
+        
         return text
     
     def _split_long_words_heuristic(self, text: str, max_length: int = 15) -> str:
-        """
-        🆕 Split suspiciously long words using heuristics
+        """ Split suspiciously long words using heuristics
         """
         words = text.split()
         result = []
@@ -1805,8 +1646,7 @@ class ArabicSpellChecker:
         return ' '.join(result)
     
     def _normalize_tanween_patterns(self, text: str) -> str:
-        """
-        🆕 Normalize tanween patterns
+        """ Normalize tanween patterns
         
         Examples:
             جدأ → جداً
@@ -1834,7 +1674,7 @@ class ArabicSpellChecker:
         text = self.postprocessor.remove_tatweel(text)
         text = self.postprocessor.normalize_special_chars(text)
         
-        # 🆕 التحسينات المدمجة (IMPROVEMENTS INTEGRATED!)
+        # Integrated improvements
         # Fix repeated chars and merged words with errors FIRST
         text = self._fix_repeated_end_chars(text)
         text = self._fix_merged_with_errors(text)
@@ -1859,38 +1699,40 @@ class ArabicSpellChecker:
         
         return text
     
-    def _fix_word_split(self, text: str) -> str:
-        """Fix over-split words by joining fragments"""
-        return self.postprocessor.join_fragments(text)
-    
     def postprocess(self, text: str, original: str = "") -> str:
-        """Postprocessing pipeline"""
-        return self.postprocessor.full_postprocess(text, original)
+        """Postprocessing pipeline — passes vocab_manager for smart ه/ة handling"""
+        return self.postprocessor.full_postprocess(text, original, vocab_manager=self.vocab_manager)
     
     def model_inference(self, text: str, num_return_sequences: int = 5) -> List[str]:
-        """Run seq2seq model inference and return top candidates"""
+        """Run seq2seq model inference and return top candidates.
+        Also extracts beam scores (token-level probabilities) for diagnostics.
+        """
         # Tokenize
         inputs = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
         
-        # Generate (encoder-decoder needs decoder_start_token_id / bos_token_id)
-        decoder_start = getattr(
-            self.model.config, 'decoder_start_token_id', None
-        ) or getattr(self.model.config, 'bos_token_id', None) or self.tokenizer.cls_token_id
-        pad_id = getattr(self.model.config, 'pad_token_id', None) or self.tokenizer.pad_token_id
+        # Generate with beam search
+        # Keeping 5 beams as model was trained/optimized for this
+        # Keeping 5 beams as model was trained/optimized for this
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,
-                max_length=128,
                 num_beams=5,
                 num_return_sequences=num_return_sequences,
                 early_stopping=True,
-                decoder_start_token_id=decoder_start,
-                pad_token_id=pad_id,
+                return_dict_in_generate=True,
+                output_scores=True
             )
         
         # Decode
-        candidates = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
+        candidates = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)
+        
+        # Store beam scores for potential use
+        self._last_beam_scores = {}
+        if hasattr(outputs, 'sequences_scores') and outputs.sequences_scores is not None:
+            scores = outputs.sequences_scores.tolist()
+            for cand, score in zip(candidates, scores):
+                self._last_beam_scores[cand] = score
         
         return candidates
     
@@ -1924,36 +1766,56 @@ class ArabicSpellChecker:
         candidates.append(preprocessed_text)
         
         # B. Smart Rules Candidate (Aggressive Heuristic)
-        # This helps when the model fails but the rule-based fix is obvious
         rules_candidate = self.rules.advanced_heuristic_repair(text)
         candidates.append(rules_candidate)
         
-        # B2. Edit Distance Candidate (NEW!)
-        # Tries to fix typos using simple edit distance (Norvig)
+        # B2. Edit Distance Candidate
         edit_candidate = self.edit_corrector.generate_candidate(text)
         if edit_candidate != text and edit_candidate != rules_candidate:
             candidates.append(edit_candidate)
         
-        # B3. Split/Merge Candidate (Phase 2)
-        # NOTE: Disabled - caused regression (Hybrid Wins 144→134)
-        # split_merge_candidate = self.split_merge.process_text(preprocessed_text)
-        # if split_merge_candidate != preprocessed_text and split_merge_candidate not in candidates:
-        #     candidates.append(split_merge_candidate)
-        
         # C. Model Beams
+        raw_model_output = None  # Track for safety net
         try:
             model_candidates = self.model_inference(preprocessed_text, num_return_sequences=5)
+            raw_model_output = model_candidates[0] if model_candidates else None
             candidates.extend(model_candidates)
             
-            # D. Word-Aligned Hybrid Candidate (Phase 2 - Solution 5)
+            # D. Word-Aligned Hybrid Candidate
             # Creates a hybrid by selecting best word from each position
-            # (OOV input + IV output → take output, IV input + OOV output → keep input)
             if model_candidates:
                 hybrid_candidate = self.word_aligner.align_words(preprocessed_text, model_candidates[0])
                 if hybrid_candidate not in candidates:
                     candidates.append(hybrid_candidate)
+                
+                # E. Word-Aligned with ALL top beams (not just beam 0)
+                for beam in model_candidates[1:3]:  # Top 3 beams
+                    hybrid_beam = self.word_aligner.align_words(preprocessed_text, beam)
+                    if hybrid_beam not in candidates:
+                        candidates.append(hybrid_beam)
+            
+            # D2. Token-level Voting Candidate
+            # Majority-vote each token across all beams
+            if model_candidates and len(model_candidates) >= 3:
+                try:
+                    beam_word_lists = [c.split() for c in model_candidates]
+                    max_words = max(len(wl) for wl in beam_word_lists)
+                    voted_words = []
+                    for pos in range(max_words):
+                        words_at_pos = []
+                        for wl in beam_word_lists:
+                            if pos < len(wl):
+                                words_at_pos.append(wl[pos])
+                        if words_at_pos:
+                            most_common = Counter(words_at_pos).most_common(1)[0][0]
+                            voted_words.append(most_common)
+                    voted_candidate = ' '.join(voted_words)
+                    if voted_candidate not in candidates:
+                        candidates.append(voted_candidate)
+                except Exception:
+                    pass
         except Exception as e:
-            print(f"⚠️ Model inference failed: {e}")
+            logger.warning(f"Model inference failed: {e}")
         
         # Remove duplicates while preserving order
         unique_candidates = []
@@ -1964,6 +1826,8 @@ class ArabicSpellChecker:
                 seen.add(c)
         candidates = unique_candidates
         
+
+        
         # 4. Rerank Candidates
         best_candidate = preprocessed_text
         best_score = -1.0
@@ -1982,7 +1846,7 @@ class ArabicSpellChecker:
                 reason = "too_short"
 
             # ═══════════════════════════════════════════════════════════════════════════
-            # NEW: VOCABULARY-AWARE ACCEPTANCE (Phase 1 - Key Fix for Raw Wins)
+            # VOCABULARY-AWARE ACCEPTANCE
             # ═══════════════════════════════════════════════════════════════════════════
             # Logic: OOV→IV = ACCEPT (boost), IV→OOV = REJECT (penalize)
             # This prevents over-conservative validation from rejecting correct corrections
@@ -2029,26 +1893,35 @@ class ArabicSpellChecker:
                 try:
                     fluency_score = self.contextual.calculate_sentence_score(cand)
                 except Exception as e:
-                    print(f"⚠️ Scoring failed: {e}")
+                    logger.warning(f"Scoring failed: {e}")
                     fluency_score = 0.5 # Default fallback
             else:
                 fluency_score = 1.0 
             
-            # C. Similarity Score (Damerau-Levenshtein Distance - Phase 1 improvement)
-            # Penalize unnecessary changes. Using DL distance: transpositions = 1 edit (not 2)
-            # This helps cases like اقصتاديا→اقتصاديا (swap صت→تص counts as 1)
+            # C. Similarity Score (Damerau-Levenshtein Distance)
             dist = VocabularyManager.damerau_levenshtein_distance(preprocessed_text, cand)
-            # Using preprocessed_text as anchor because it has basic normalization.
-            # Comparison with 'original' might penalize fixing harakat/spelling.
-            
             max_len = max(len(preprocessed_text), len(cand), 1)
             similarity = 1.0 - (dist / max_len)
             
-            # Boost matches
+            # Boost exact matches
             if cand == preprocessed_text:
                 similarity = 1.0
             
-            # NEW: HIGH CONFIDENCE GATING (Phase 1/3 - Solution)
+            # Keyboard Proximity Bonus
+            # If changes between input and candidate are keyboard-adjacent,
+            # it's more likely a typo fix (give bonus)
+            keyboard_bonus = 1.0
+            input_words = preprocessed_text.split()
+            cand_words = cand.split()
+            if len(input_words) == len(cand_words):
+                for iw, cw in zip(input_words, cand_words):
+                    if iw != cw and len(iw) == len(cw):
+                        # Check char-by-char differences
+                        for ic, cc in zip(iw, cw):
+                            if ic != cc and RulesBasedCorrector.is_keyboard_neighbor(ic, cc):
+                                keyboard_bonus *= 1.05  # 5% bonus per keyboard-adjacent fix
+            
+            # HIGH CONFIDENCE GATING
             # If model is extremely confident (high fluency) and words are valid, relax validation
             # This allows correcting severe corruptions that fail strict edit distance
             if fluency_score > 0.85 and cand_oov_count == 0:
@@ -2060,24 +1933,16 @@ class ArabicSpellChecker:
                           vocab_boost *= 1.2  # Bonus for high confidence
                           validity_factor = 1.0  # Reset validity factor
             
-            # Final Score
-            # Fluency is roughly [0, 1] (prob). Similarity [0, 1].
-            # We want to balance staying close to original vs being fluent.
-            # If fluency is very low, it's garbage.
-            # If similarity is very low, it's hallucination.
-            
-            # Weighting:
-            # We value Similarity highly to be conservative.
-            # But we need Fluency to break ties or fix errors.
+            # Final Score = (Fluency^0.3) * (Similarity^3.0) * Validity * VocabBoost * KeyboardBonus * BeamBoost
+            fluency_exp = 0.3
+            similarity_exp = 3.0
             
-            # New Formula:
-            # Score = (Fluency^0.3) * (Similarity^2.0) * Validity * VocabBoost
-            # Using exponent to control trade-off. 
-            # High sim power -> prefers closer matches.
-            # Low fluency power -> flattens probability differences (since probs are small).
-            # VocabBoost: rewards OOV→IV, penalizes IV→OOV
+            # Beam 0 Boost — model's top beam gets 15% priority
+            beam_boost = 1.0
+            if raw_model_output and cand == raw_model_output:
+                beam_boost = 1.15
             
-            final_score = (fluency_score ** 0.3) * (similarity ** 3.0) * validity_factor * vocab_boost
+            final_score = (fluency_score ** fluency_exp) * (similarity ** similarity_exp) * validity_factor * vocab_boost * keyboard_bonus * beam_boost
             
             candidate_scores.append({
                 'text': cand,
@@ -2085,7 +1950,7 @@ class ArabicSpellChecker:
                 'reason': reason,
                 'fluency': fluency_score,
                 'similarity': similarity,
-                'vocab_boost': vocab_boost,  # NEW: Track vocab boost
+                'vocab_boost': vocab_boost,
                 'input_oov': input_oov_count,
                 'cand_oov': cand_oov_count,
                 'final_score': final_score
@@ -2095,32 +1960,265 @@ class ArabicSpellChecker:
                 best_score = final_score
                 best_candidate = cand
         
+        # ═══════════════════════════════════════════════════════════════════════════
+        # --- Output Quality Scoring (Minimum Score Threshold) ---
+        # If ALL candidates scored poorly, the correction is unreliable → keep input
+        # ═══════════════════════════════════════════════════════════════════════════
+        if best_candidate != preprocessed_text:
+            # Check: did the best candidate actually get a decent score?
+            # The preprocessed input (candidate 0) is always in the pool.
+            # If the best candidate barely beats preprocessed_text, it might not be trustworthy.
+            preprocessed_score = 0.0
+            for cs in candidate_scores:
+                if cs['text'] == preprocessed_text:
+                    preprocessed_score = cs['final_score']
+                    break
+            
+            # If best score is less than 1.05x the preprocessed score AND
+            # the best candidate introduced OOV words → fall back to preprocessed
+            if preprocessed_score > 0 and best_score < preprocessed_score * 1.05:
+                best_oov = self.vocab_manager.count_oov_words(best_candidate)
+                prep_oov = self.vocab_manager.count_oov_words(preprocessed_text)
+                if best_oov > prep_oov:
+                    best_candidate = preprocessed_text
+                    best_score = preprocessed_score
+        
+        # ═══════════════════════════════════════════════════════════════════════════
+        # --- Contextual Validation Layer ---
+        # Compare fluency of input vs best candidate
+        # If correction made text LESS fluent → reject the correction
+        # ═══════════════════════════════════════════════════════════════════════════
+        if best_candidate != preprocessed_text and self.use_contextual and self.contextual:
+            try:
+                input_fluency = self.contextual.calculate_sentence_score(preprocessed_text)
+                best_fluency = 0.0
+                for cs in candidate_scores:
+                    if cs['text'] == best_candidate:
+                        best_fluency = cs['fluency']
+                        break
+                
+                # If input is significantly more fluent than best candidate
+                # AND both have similar OOV counts → prefer input
+                if input_fluency > 0 and best_fluency > 0:
+                    if input_fluency > best_fluency * 1.5:  # Input 50% more fluent
+                        input_oov = self.vocab_manager.count_oov_words(preprocessed_text)
+                        best_oov = self.vocab_manager.count_oov_words(best_candidate)
+                        if input_oov <= best_oov:
+                            # Input is more fluent AND has fewer/equal OOV → keep input
+                            best_candidate = preprocessed_text
+            except Exception:
+                pass  # Contextual validation is optional
+        
         # 5. Postprocess Winner
         result = self.postprocess(best_candidate, original)
         
+        # 5.5 IV-Safe Postprocessing Check
+        # If postprocessing changed an IV word to OOV, revert that word
+        if result != best_candidate:
+            result_words = result.split()
+            best_words = best_candidate.split()
+            if len(result_words) == len(best_words):
+                fixed_words = []
+                input_words_pp = preprocessed_text.split()
+                for idx_fw, (rw, bw) in enumerate(zip(result_words, best_words)):
+                    if rw != bw:
+                        # Postprocessor changed this word
+                        bw_iv = self.vocab_manager.is_iv(bw)
+                        rw_iv = self.vocab_manager.is_iv(rw)
+                        if bw_iv and not rw_iv:
+                            # IV → OOV: revert to pre-postprocess version
+                            fixed_words.append(bw)
+                        elif bw_iv and rw_iv:
+                            # Postprocess Distance Guard
+                            # DISABLED: Caused word-level regression. When both are IV,
+                            # the postprocessor's choice (rw) is usually better because
+                            # it applies Arabic-specific rules (hamza, ta marbuta).
+                            fixed_words.append(rw)
+                        else:
+                            fixed_words.append(rw)
+                    else:
+                        fixed_words.append(rw)
+                result = ' '.join(fixed_words)
+        
         # 6. Contextual fine-tuning (BERT Masked Refinement)
-        # Note: Applying to full sentence (OOV-only mode caused regression due to lack of context)
+        # IV-Safe mode - pass vocab_manager to protect IV words
+        # BERT Kill Switch - also pass raw_model_output to protect model-confident words
         if self.use_contextual and self.contextual:
              if len(result) > 3:
-                 result = self.contextual.refine_sentence_with_mask(result)
+                 result = self.contextual.refine_sentence_with_mask(
+                     result, vocab_manager=self.vocab_manager,
+                     raw_model_output=raw_model_output
+                 )
         
-        # 7. Phase 2: Safe Split/Merge Post-processing
+        # 7. Safe Split/Merge Post-processing
         # Only apply merge_fragments (safe: only merges when result is IV)
-        # This fixes ta-marbuta detachment like السوري ة → السورية
         result = self.split_merge.merge_fragments(result)
         
+        # ═══════════════════════════════════════════════════════════════════════════
+        # VALIDATION & QUALITY CHECKS
+        # ═══════════════════════════════════════════════════════════════════════════
+        
+        # 8. Output Stability Test (Solution 30)
+        # If correcting the output again changes it → unstable correction → reject
+        # Stable corrections are idempotent: correct(correct(x)) == correct(x)
+        if result != preprocessed_text and raw_model_output:
+            try:
+                # Quick stability check: run the result through preprocessing only
+                # (full model inference would be too slow)
+                re_preprocessed = self.preprocess(result)
+                
+                # If re-preprocessing changes the result significantly, it was unstable
+                stability_dist = VocabularyManager.damerau_levenshtein_distance(result, re_preprocessed)
+                result_len = max(len(result), 1)
+                
+                if stability_dist > 0:
+                    # Result is not stable under re-preprocessing
+                    stability_ratio = stability_dist / result_len
+                    
+                    if stability_ratio > 0.15:  # More than 15% changed → very unstable
+                        # Fall back to raw model output if it's more stable
+                        raw_re = self.preprocess(raw_model_output)
+                        raw_stability = VocabularyManager.damerau_levenshtein_distance(
+                            raw_model_output, raw_re
+                        ) / max(len(raw_model_output), 1)
+                        
+                        if raw_stability < stability_ratio:
+                            # Raw is more stable → use it
+                            raw_oov = self.vocab_manager.count_oov_words(raw_model_output)
+                            our_oov = self.vocab_manager.count_oov_words(result)
+                            if raw_oov <= our_oov:
+                                result = raw_model_output
+            except Exception:
+                pass  # Stability check is optional, don't break pipeline
+        
+        # 9. Bidirectional Word-Level Validation (Solution 24)
+        # Compare our result word-by-word with raw model output
+        # If we corrupted a word that the model got right, revert that word
+        if raw_model_output and result != raw_model_output:
+            result_words = result.split()
+            raw_words = raw_model_output.split()
+            
+            if len(result_words) == len(raw_words):
+                corrected_words = []
+                changed = False
+                
+                for rw, raw_w in zip(result_words, raw_words):
+                    if rw != raw_w:
+                        rw_iv = self.vocab_manager.is_iv(rw)
+                        raw_iv = self.vocab_manager.is_iv(raw_w)
+                        
+                        # Case 1: Our word is OOV but raw word is IV → take raw
+                        if not rw_iv and raw_iv:
+                            corrected_words.append(raw_w)
+                            changed = True
+                        # Case 2: Both IV but our word is further from input
+                        elif rw_iv and raw_iv:
+                            # Find corresponding input word
+                            input_words = preprocessed_text.split()
+                            idx = len(corrected_words)
+                            if idx < len(input_words):
+                                input_w = input_words[idx]
+                                rw_dist = Levenshtein.distance(input_w, rw)
+                                raw_dist = Levenshtein.distance(input_w, raw_w)
+                                # If raw is closer to input AND both are IV → prefer raw
+                                # (our pipeline likely introduced unnecessary change)
+                                if raw_dist < rw_dist:
+                                    corrected_words.append(raw_w)
+                                    changed = True
+                                else:
+                                    corrected_words.append(rw)
+                            else:
+                                corrected_words.append(rw)
+                        else:
+                            corrected_words.append(rw)
+                    else:
+                        corrected_words.append(rw)
+                
+                if changed:
+                    new_result = ' '.join(corrected_words)
+                    # Only accept if the new result doesn't increase OOV
+                    new_oov = self.vocab_manager.count_oov_words(new_result)
+                    old_oov = self.vocab_manager.count_oov_words(result)
+                    if new_oov <= old_oov:
+                        result = new_result
+        
+        # 10. SAFETY NET: Compare with raw model output (Conservative)
+        # Only switch to raw if raw is CLEARLY better
+        if raw_model_output and raw_model_output != result:
+            raw_oov = self.vocab_manager.count_oov_words(raw_model_output)
+            our_oov = self.vocab_manager.count_oov_words(result)
+            
+            # Case A: Raw all-IV, ours has OOV
+            if raw_oov == 0 and our_oov > 0:
+                is_valid, reason = self.validator.validate(original, raw_model_output, "mixed")
+                if is_valid or reason == "space_leniency_accept":
+                    result = raw_model_output
+            
+            # Case B: Both all-IV but raw is more similar to input
+            # Catches BERT/postprocess damage (word substitutions up to 5 char distance)
+            elif raw_oov == 0 and our_oov == 0:
+                raw_dist = VocabularyManager.damerau_levenshtein_distance(original, raw_model_output)
+                our_dist = VocabularyManager.damerau_levenshtein_distance(original, result)
+                result_vs_raw_dist = VocabularyManager.damerau_levenshtein_distance(result, raw_model_output)
+                # Threshold at 3 chars — covers single char edits and small substitutions
+                # (widening to 5 caused regression by reverting valid hybrid corrections)
+                if raw_dist < our_dist and result_vs_raw_dist <= 3:
+                    raw_valid, _ = self.validator.validate(original, raw_model_output, "mixed")
+                    if raw_valid:
+                        result = raw_model_output
+            
+            # Case C: Word count differs — raw might have correct splitting
+            # Catches: 'فيلق → في فيلق' (pipeline added word)
+            # or 'بلاكبيرن روفرز → بلاكبيرن روفر' (pipeline lost word ending)
+            elif raw_oov == 0:
+                raw_wc = len(raw_model_output.split())
+                our_wc = len(result.split())
+                if raw_wc != our_wc:
+                    raw_dist = VocabularyManager.damerau_levenshtein_distance(original, raw_model_output)
+                    our_dist = VocabularyManager.damerau_levenshtein_distance(original, result)
+                    if raw_dist < our_dist:
+                        raw_valid, _ = self.validator.validate(original, raw_model_output, "mixed")
+                        if raw_valid:
+                            result = raw_model_output
+        
         return result
 
+# ═══════════════════════════════════════════════════════════════════════════════
+# PUBLIC API
+# ═══════════════════════════════════════════════════════════════════════════════
+
+# Exported for use by benchmark.py and external consumers
+spell_checker = None  # Will be initialized on first import with __main__ or by benchmark
+
+
+def initialize(use_contextual=True):
+    """Initialize the spell checker. Call once before using."""
+    global spell_checker
+    spell_checker = ArabicSpellChecker(model, tokenizer, device, use_contextual=use_contextual)
+    logger.info("Spell checker initialized")
+    return spell_checker
+
+
+if __name__ == "__main__":
+    sc = initialize(use_contextual=True)
+
+    # Quick demo
+    test_cases = [
+        "السلام عليكممم",
+        "فيالمدرسه",
+        "الطقص جميل اليومم",
+    ]
 
-print("✅ All classes defined successfully!")
-print("   - ErrorType")
-print("   - AraSpellPostProcessor")
-print("   - ErrorClassifier")
-print("   - RulesBasedCorrector")
-print("   - OutputValidator")
-print("   - ContextualCorrector")
-print("   - ArabicSpellChecker")
+    print("\n" + "=" * 60)
+    print("AraSpell Demo")
+    print("=" * 60)
 
+    for text in test_cases:
+        corrected = sc.correct(text)
+        print(f"\n  Input:     {text}")
+        print(f"  Corrected: {corrected}")
 
-print("✅ AraSpell classes loaded successfully")
+    print("\n" + "=" * 60)
+    print("For full benchmark, run: python benchmark.py")
+    print("=" * 60)
 

Dockerfile

@@ -13,19 +13,34 @@ COPY requirements.txt .
 RUN pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu && \
     pip install --no-cache-dir -r requirements.txt
 
-# Pre-download the summarization model during build (network is available here)
+# Pre-download models during build (network is available here)
 # At runtime, the container has NO outbound DNS, so models must be cached
+
+# 1. Summarization model (MBart, float16)
 RUN python -c "\
 from transformers import MBartForConditionalGeneration, AutoTokenizer, AutoConfig; \
 import torch; \
 repo = 'bayan10/summarization-model'; \
-print('Downloading tokenizer...'); \
+print('Downloading summarization tokenizer...'); \
 AutoTokenizer.from_pretrained(repo); \
-print('Downloading config...'); \
+print('Downloading summarization config...'); \
 AutoConfig.from_pretrained(repo); \
-print('Downloading model (float16)...'); \
+print('Downloading summarization model (float16)...'); \
 MBartForConditionalGeneration.from_pretrained(repo, torch_dtype=torch.float16); \
-print('Model cached successfully!'); \
+print('Summarization model cached!'); \
+"
+
+# 2. Spelling model (AraSpell — AraBERT encoder-decoder + checkpoint)
+RUN python -c "\
+from huggingface_hub import hf_hub_download; \
+from transformers import AutoTokenizer, EncoderDecoderModel; \
+print('Downloading AraSpell checkpoint...'); \
+hf_hub_download(repo_id='bayan10/AraSpell-Model', filename='last_model.pt'); \
+print('Downloading AraBERT tokenizer...'); \
+AutoTokenizer.from_pretrained('aubmindlab/bert-base-arabertv02'); \
+print('Downloading AraBERT encoder-decoder...'); \
+EncoderDecoderModel.from_encoder_decoder_pretrained('aubmindlab/bert-base-arabertv02', 'aubmindlab/bert-base-arabertv02'); \
+print('Spelling model cached!'); \
 "
 
 # Copy application code

gradio Spelling.py

@@ -0,0 +1,210 @@
+import gradio as gr
+
+import re
+from AraSpell import initialize
+
+# تهيئة المصحح الإملائي وتحميل الموديل
+sc = initialize(use_contextual=True)
+
+import Levenshtein
+
+def align_words(in_words, out_words):
+    n = len(in_words)
+    m = len(out_words)
+    dp = [[0] * (m + 1) for _ in range(n + 1)]
+    
+    for i in range(1, n + 1):
+        dp[i][0] = dp[i-1][0] + len(in_words[i-1])
+    for j in range(1, m + 1):
+        dp[0][j] = dp[0][j-1] + len(out_words[j-1])
+        
+    for i in range(1, n + 1):
+        for j in range(1, m + 1):
+            cost_replace = Levenshtein.distance(in_words[i-1], out_words[j-1])
+            dp[i][j] = min(
+                dp[i-1][j-1] + cost_replace,
+                dp[i-1][j] + len(in_words[i-1]),
+                dp[i][j-1] + len(out_words[j-1])
+            )
+            
+    i, j = n, m
+    ops = []
+    while i > 0 or j > 0:
+        if i > 0 and j > 0:
+            cost_replace = Levenshtein.distance(in_words[i-1], out_words[j-1])
+            if dp[i][j] == dp[i-1][j-1] + cost_replace:
+                if cost_replace == 0:
+                    ops.append(('equal', [in_words[i-1]], [out_words[j-1]]))
+                else:
+                    ops.append(('replace', [in_words[i-1]], [out_words[j-1]]))
+                i -= 1
+                j -= 1
+                continue
+        if i > 0 and dp[i][j] == dp[i-1][j] + len(in_words[i-1]):
+            ops.append(('delete', [in_words[i-1]], []))
+            i -= 1
+        else:
+            ops.append(('insert', [], [out_words[j-1]]))
+            j -= 1
+            
+    ops.reverse()
+    
+    merged_ops = []
+    for op, in_w, out_w in ops:
+        if not merged_ops:
+            merged_ops.append([op, in_w, out_w])
+            continue
+            
+        prev_op = merged_ops[-1][0]
+        if op != 'equal' and prev_op != 'equal' and not (op == 'replace' and prev_op == 'replace'):
+            merged_ops[-1][0] = 'replace'
+            merged_ops[-1][1].extend(in_w)
+            merged_ops[-1][2].extend(out_w)
+        else:
+            merged_ops.append([op, in_w, out_w])
+            
+    return merged_ops
+
+def generate_highlights(input_text):
+    if not input_text or not input_text.strip():
+        return [], {}
+        
+    corrected_text = sc.correct(input_text)
+    
+    in_words = input_text.split()
+    out_words = corrected_text.split()
+    
+    ops = align_words(in_words, out_words)
+    
+    highlight_list = []
+    suggestions_map = {}
+    
+    idx = 0
+    for tag, in_w, out_w in ops:
+        if tag == 'equal':
+            for w in in_w:
+                highlight_list.append((w, None))
+                highlight_list.append((" ", None))
+                idx += 2
+        elif tag == 'replace' or tag == 'insert' or tag == 'delete':
+            in_phrase = " ".join(in_w) if in_w else "[ناقص]"
+            out_phrase = " ".join(out_w) if out_w else "(حذف الكلمة)"
+            
+            highlight_list.append((in_phrase, " "))
+            sugs = [out_phrase]
+            
+            if len(in_w) == 1 and len(out_w) == 1:
+                clean_w = re.sub(r'[^\w]', '', in_w[0])
+                try:
+                    edit_cands = sc.edit_corrector.known(sc.edit_corrector.edits1(clean_w))
+                    if edit_cands:
+                        edit_cands = sorted(list(edit_cands), key=lambda x: sc.vocab_manager.get_frequency_rank(x))
+                        for c in edit_cands:
+                            if c not in sugs and len(sugs) < 3:
+                                sugs.append(c)
+                except Exception:
+                    pass
+                    
+            suggestions_map[idx] = sugs
+            highlight_list.append((" ", None))
+            idx += 2
+            
+    if highlight_list and highlight_list[-1] == (" ", None):
+        highlight_list.pop()
+        
+    return highlight_list, suggestions_map
+
+# ==========================================
+# تصميم واجهة المستخدم التفاعلية (Gradio Blocks)
+# ==========================================
+with gr.Blocks(theme=gr.themes.Soft(), css="""
+    .highlight-error { background-color: #ffcccc !important; border-radius: 4px; padding: 2px; }
+    .rtl-text { direction: rtl !important; text-align: right !important; }
+""") as iface:
+    
+    gr.Markdown("# 📝 AraSpell - المصحح الإملائي التفاعلي")
+    gr.Markdown("أدخل النص أدناه واضغط على **فحص النص**. سيقوم النظام بتلوين الأخطاء باللون الأحمر. **انقر على الكلمة الملونة** لتظهر لك خيارات التصحيح أسفلها!")
+    
+    # متغيرات حالة (State) لحفظ البيانات خلف الكواليس
+    suggestions_state = gr.State({})
+    current_edit_index = gr.State(None)
+    highlight_list_state = gr.State([])
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_box = gr.Textbox(lines=8, label="النص الأصلي", placeholder="أدخل النص العربي هنا...")
+            check_btn = gr.Button("🔍 فحص النص", variant="primary")
+            
+        with gr.Column(scale=1):
+            output_highlights = gr.HighlightedText(
+                label="النتيجة (اضغط على الكلمات الملونة للتصحيح)",
+                combine_adjacent=False,
+                show_legend=False,
+                color_map={" ": "red"},
+                elem_classes="rtl-text"
+            )
+            
+            # لوحة الاقتراحات (مخفية في البداية)
+            with gr.Group(visible=False) as suggestion_panel:
+                gr.Markdown("### 💡 اختر التصحيح المناسب:")
+                suggestion_radio = gr.Radio(choices=[], label="")
+                apply_btn = gr.Button("✅ تطبيق التصحيح", variant="secondary")
+
+    # 1. عند الضغط على فحص النص
+    def process_text(text):
+        h_list, s_map = generate_highlights(text)
+        # إرجاع: النص المظلل، قاموس الاقتراحات، قائمة التظليل (State)، وإخفاء لوحة الاقتراحات
+        return h_list, s_map, h_list, gr.update(visible=False)
+
+    check_btn.click(
+        fn=process_text,
+        inputs=[input_box],
+        outputs=[output_highlights, suggestions_state, highlight_list_state, suggestion_panel]
+    )
+    
+    # 2. عند النقر على أي كلمة داخل النص المظلل
+    def on_highlight_click(evt: gr.SelectData, s_map):
+        index = evt.index
+        # معالجة مشكلة تحويل المفاتيح إلى نصوص (Strings) في Gradio State
+        if index in s_map:
+            choices = s_map[index]
+        elif str(index) in s_map:
+            choices = s_map[str(index)]
+        else:
+            # إخفاء اللوحة إذا ضغط على كلمة صحيحة
+            return gr.update(visible=False), gr.update(), None
+            
+        # إظهار اللوحة وتحديث الخيارات
+        return gr.update(visible=True), gr.update(choices=choices, value=choices[0]), index
+
+    output_highlights.select(
+        fn=on_highlight_click,
+        inputs=[suggestions_state],
+        outputs=[suggestion_panel, suggestion_radio, current_edit_index]
+    )
+    
+    # 3. عند اختيار اقتراح والضغط على "تطبيق"
+    def apply_correction(choice, edit_idx, h_list):
+        if edit_idx is not None and choice:
+            # تحديث الكلمة في قائمة التظليل (بدون إعادة تشغيل الموديل لتكون سريعة جداً)
+            if choice == "(حذف الكلمة)":
+                h_list[edit_idx] = ("", None)
+            else:
+                h_list[edit_idx] = (choice, None)
+                
+            # إعادة بناء النص الجديد
+            new_text = "".join([t[0] for t in h_list])
+            
+            # إرجاع: تحديث مربع الإدخال، التظليل الجديد، بقاء الاقتراحات كما هي، State الجديد، وإخفاء اللوحة
+            return new_text, h_list, gr.update(), h_list, gr.update(visible=False)
+            
+        return gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+
+    apply_btn.click(
+        fn=apply_correction,
+        inputs=[suggestion_radio, current_edit_index, highlight_list_state],
+        outputs=[input_box, output_highlights, suggestions_state, highlight_list_state, suggestion_panel]
+    )
+
+if __name__ == "__main__":
+    iface.launch()

src/app.py

@@ -153,7 +153,7 @@ def health_check():
             'mode': 'hf_spaces_local',
             'models': {
                 'summarization': summarization_model is not None,
-                'spelling': False,
+                'spelling': _spelling_available(),
                 'autocomplete': False,
                 'grammar': False,
                 'punctuation': False
@@ -223,6 +223,74 @@ def debug_models():
     }), 200
 
 
+def _spelling_available():
+    """Check if spelling model is loaded (without triggering lazy load)."""
+    try:
+        from nlp.spelling.araspell_service import is_loaded
+        return is_loaded()
+    except Exception:
+        return False
+
+
+@app.route('/api/spelling', methods=['POST'])
+def spelling_correction():
+    """
+    Correct spelling in Arabic text.
+    
+    Request JSON:
+    {
+        "text": "Arabic text with spelling errors"
+    }
+    
+    Response JSON:
+    {
+        "original_text": "...",
+        "corrected_text": "...",
+        "status": "success"
+    }
+    """
+    try:
+        if not request.is_json:
+            return jsonify({'error': 'Request must be JSON', 'status': 'error'}), 400
+        
+        data = request.get_json()
+        text = data.get('text', '').strip()
+        
+        if not text:
+            return jsonify({'error': 'Text is required', 'status': 'error'}), 400
+        
+        if len(text) > MAX_TEXT_LENGTH:
+            return jsonify({
+                'error': f'Text too long. Maximum {MAX_TEXT_LENGTH} characters.',
+                'status': 'error'
+            }), 400
+        
+        logger.info(f"Spelling correction request: text_length={len(text)}")
+        
+        from nlp.spelling.araspell_service import get_spelling_model
+        checker = get_spelling_model()
+        corrected = checker.correct(text)
+        
+        return jsonify({
+            'original_text': text,
+            'corrected_text': corrected,
+            'status': 'success'
+        }), 200
+    
+    except RuntimeError as e:
+        logger.error(f"Spelling model error: {e}")
+        return jsonify({
+            'error': f'Spelling model unavailable: {str(e)[:200]}',
+            'status': 'error'
+        }), 503
+    except Exception as e:
+        logger.error(f"Spelling correction error: {e}")
+        return jsonify({
+            'error': f'Spelling correction failed: {str(e)[:200]}',
+            'status': 'error'
+        }), 500
+
+
 @app.route('/api/summarize', methods=['POST'])
 def summarize():
     """

src/nlp/__init__.py

@@ -0,0 +1 @@
+# AraSpell — Arabic Spelling Correction Module

src/nlp/spelling/__init__.py

@@ -0,0 +1 @@
+# AraSpell — Arabic Spelling Correction

src/nlp/spelling/araspell_rules.py

@@ -0,0 +1,1407 @@
+# AraSpell — Arabic Spell Checker Pipeline (Rules & Classes)
+# Extracted from AraSpell.py — NO global model loading, NO Gradio dependencies.
+# All classes are imported by araspell_service.py.
+
+import re
+import math
+import logging
+import torch
+from collections import Counter
+from enum import Enum
+from typing import List, Tuple, Optional
+
+import Levenshtein
+import jellyfish
+
+logger = logging.getLogger(__name__)
+
+# ─────────────────────────────────────────────────────────────────────────────
+# ERROR TYPE ENUM
+# ─────────────────────────────────────────────────────────────────────────────
+
+class ErrorType(Enum):
+    """Types of spelling errors"""
+    CHAR_REPETITION = "char_repetition"
+    WORD_MERGE = "word_merge"
+    CHAR_SUBSTITUTION = "char_substitution"
+    MIXED = "mixed"
+    CLEAN = "clean"
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# POST PROCESSOR
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class AraSpellPostProcessor:
+    """Arabic text post-processing techniques."""
+    
+    ARABIC_HARAKAT = 'ًٌٍَُِّْ'
+    TATWEEL = 'ـ'
+    NORMALIZER_MAP = {
+        'ﻹ': 'لإ', 'ﻷ': 'لأ', 'ﻵ': 'لآ', 'ﻻ': 'لا', 'ﷲ': 'الله'
+    }
+    ARABIC_CONSONANTS = set('بتثجحخدذرزسشصضطظعغفقكلمن')
+    
+    # --- Basic Normalization ---
+    
+    @staticmethod
+    def remove_harakat(text: str) -> str:
+        """Remove Arabic diacritics"""
+        return re.sub(r'[ً-ْ]', '', text)
+    
+    @staticmethod
+    def remove_tatweel(text: str) -> str:
+        """Remove Arabic kashida/tatweel"""
+        return text.replace(AraSpellPostProcessor.TATWEEL, '')
+    
+    @staticmethod
+    def normalize_special_chars(text: str) -> str:
+        """Normalize special Arabic ligatures"""
+        for old, new in AraSpellPostProcessor.NORMALIZER_MAP.items():
+            text = text.replace(old, new)
+        return text
+    
+    # --- Core Functions ---
+    
+    @staticmethod
+    def unified_collapse_repeated(text: str) -> str:
+        """
+        Collapse repeated characters.
+        Arabic: 3+ consecutive → 1 | Latin: 2+ consecutive → 1
+        """
+        text = re.sub(r"([\u0600-\u06FF])\1{2,}", r"\1", text)
+        text = re.sub(r"([a-zA-Z])\1+", r"\1", text)
+        return text
+    
+    @staticmethod
+    def remove_duplicate_words(text: str) -> str:
+        """Remove consecutive duplicate words. e.g. كتاب كتاب → كتاب"""
+        words = text.split()
+        if len(words) < 2:
+            return text
+        result = [words[0]]
+        for i in range(1, len(words)):
+            if words[i] != words[i-1]:
+                result.append(words[i])
+        return ' '.join(result)
+    
+    @staticmethod
+    def normalize_spaces(text: str) -> str:
+        """Normalize whitespace: multiple spaces, unicode spaces, punctuation spacing."""
+        text = re.sub(r' +', ' ', text)
+        text = text.replace('\u00A0', ' ')
+        text = text.replace('\u200B', '')
+        text = text.replace('\u200C', '')
+        text = text.replace('\u200D', '')
+        text = text.strip()
+        text = re.sub(r'\s*([،؛؟!.])\s*', r'\1 ', text)
+        text = text.strip()
+        return text
+    
+    @staticmethod
+    def remove_word_repetition_with_wa(text: str) -> str:
+        """Remove word و word → word"""
+        words = text.split()
+        result = []
+        i = 0
+        while i < len(words):
+            if i + 2 < len(words) and words[i] == words[i+2] and words[i+1] == 'و':
+                result.append(words[i])
+                i += 3
+            else:
+                result.append(words[i])
+                i += 1
+        return ' '.join(result)
+    
+    # --- Hamza & Ta Marbuta Handling ---
+    
+    @staticmethod
+    def fix_hamza_conservative(text: str) -> str:
+        """Conservative Hamza normalization — only at word END, not middle."""
+        words = text.split()
+        result = []
+        for word in words:
+            if len(word) >= 3:
+                if word.endswith('أ'):
+                    word = word[:-1] + 'ا'
+                if word.endswith('إ'):
+                    word = word[:-1] + 'ا'
+            result.append(word)
+        return ' '.join(result)
+    
+    @staticmethod
+    def fix_ha_ta_marbuta(text: str, vocab_manager=None) -> str:
+        """
+        Smart ه → ة fix at end of words.
+        Strategy: Only convert if the ة version is IV (in tokenizer vocab).
+        """
+        PROTECTED_ENDINGS = ['لله']
+        words = text.split()
+        result = []
+        for word in words:
+            if any(word.endswith(e) for e in PROTECTED_ENDINGS):
+                result.append(word)
+                continue
+            if len(word) >= 4 and word.endswith('ه'):
+                if word[-2] in AraSpellPostProcessor.ARABIC_CONSONANTS:
+                    candidate_with_ta = word[:-1] + 'ة'
+                    if vocab_manager:
+                        ta_iv = vocab_manager.is_iv(candidate_with_ta)
+                        ha_iv = vocab_manager.is_iv(word)
+                        if ta_iv:
+                            result.append(candidate_with_ta)
+                            continue
+                        elif ha_iv:
+                            result.append(word)
+                            continue
+                    else:
+                        result.append(candidate_with_ta)
+                        continue
+            result.append(word)
+        return ' '.join(result)
+    
+    # --- Hallucination Removal ---
+    
+    @staticmethod
+    def remove_hallucinations(text: str) -> str:
+        """Remove model hallucinations: duplicate words, trailing 'و' artifacts."""
+        words = text.split()
+        if not words:
+            return text
+        result = []
+        i = 0
+        
+        def normalize_word(w: str) -> str:
+            w = w.replace('ال', '').replace('ة', 'ه')
+            w = re.sub(r'[أإآ]', 'ا', w)
+            return w
+        
+        while i < len(words):
+            word = words[i]
+            if len(word) > 4 and word.endswith('و'):
+                prev_char = word[-2]
+                if prev_char in 'ةهاأإآء':
+                    word = word[:-1]
+            if i + 1 < len(words):
+                next_word = words[i + 1]
+                if normalize_word(word) == normalize_word(next_word):
+                    keep = next_word if next_word.startswith('ال') and not word.startswith('ال') else word
+                    result.append(keep)
+                    i += 2
+                    continue
+            result.append(word)
+            i += 1
+        return ' '.join(result)
+    
+    @staticmethod
+    def remove_hallucinated_prefix(text: str, original: str) -> str:
+        """Remove particles (و/في) added by model if not in original"""
+        if not original:
+            return text
+        if text.startswith('و ') and not original.startswith('و'):
+            rest = text[2:].strip()
+            if AraSpellPostProcessor.normalize_special_chars(rest) == AraSpellPostProcessor.normalize_special_chars(original):
+                return rest
+        return text
+    
+    # --- Word Splitting & Merging ---
+    
+    @staticmethod
+    def merge_separated_al(text: str) -> str:
+        """Merge 'ال' separated by space: ال + كتاب → الكتاب"""
+        return re.sub(r'\bال\s+(\w+)', r'ال\1', text)
+    
+    @staticmethod
+    def join_fragments(text: str) -> str:
+        """Join short fragments with validation."""
+        words = text.split()
+        if len(words) < 2:
+            return text
+        STANDALONE_WORDS = {
+            'من', 'في', 'على', 'عن', 'مع', 'إلى', 'الى', 'حتى', 'منذ', 'خلال', 
+            'بعد', 'قبل', 'ب', 'ل', 'ك', 'و', 'أو', 'لا', 'ما', 'لم', 'لن',
+            'هو', 'هي', 'هم', 'أن', 'إن', 'كل', 'كان', 'قد', 'قال', 'ذلك',
+            'هذا', 'هذه', 'تلك', 'التي', 'الذي', 'التى', 'اللذي'
+        }
+        result = []
+        i = 0
+        while i < len(words):
+            word = words[i]
+            if i + 1 < len(words):
+                next_word = words[i + 1]
+                if word in STANDALONE_WORDS and next_word in STANDALONE_WORDS:
+                    result.append(word)
+                    i += 1
+                    continue
+                if len(next_word) == 1:
+                    result.append(word + next_word)
+                    i += 2
+                    continue
+                if len(word) >= 2 and len(next_word) >= 2 and word[-1] == next_word[0]:
+                    if not (word in STANDALONE_WORDS and next_word in STANDALONE_WORDS):
+                        result.append(word[:-1] + next_word)
+                        i += 2
+                        continue
+                if (2 <= len(word) <= 4 and 
+                    1 <= len(next_word) <= 2 and
+                    3 <= len(word) + len(next_word) <= 7):
+                    if not (word in STANDALONE_WORDS and next_word in STANDALONE_WORDS):
+                        result.append(word + next_word)
+                        i += 2
+                        continue
+            result.append(word)
+            i += 1
+        return ' '.join(result)
+    
+    # --- Main Pipelines ---
+    
+    @staticmethod
+    def full_postprocess(text: str, original: str = "", vocab_manager=None) -> str:
+        """Apply all post-processing steps."""
+        if original:
+            text = AraSpellPostProcessor.remove_hallucinated_prefix(text, original)
+        text = AraSpellPostProcessor.normalize_special_chars(text)
+        text = AraSpellPostProcessor.remove_hallucinations(text)
+        text = AraSpellPostProcessor.unified_collapse_repeated(text)
+        text = AraSpellPostProcessor.fix_hamza_conservative(text)
+        text = AraSpellPostProcessor.fix_ha_ta_marbuta(text, vocab_manager=vocab_manager)
+        text = AraSpellPostProcessor.remove_word_repetition_with_wa(text)
+        text = AraSpellPostProcessor.remove_duplicate_words(text)
+        text = AraSpellPostProcessor.normalize_spaces(text)
+        return text
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# ERROR CLASSIFIER
+# ─────────────────────────────────────────────────────────────────────────────
+
+class ErrorClassifier:
+    """Classify type of spelling error"""
+    
+    NON_ARABIC_KEYBOARD = set('پگچژکەڕڤڵڎےۀۃھیټډڼڑ')
+    
+    @staticmethod
+    def has_char_substitution(text: str) -> bool:
+        return any(c in ErrorClassifier.NON_ARABIC_KEYBOARD for c in text)
+    
+    @staticmethod
+    def has_char_repetition(text: str, threshold: int = 3) -> bool:
+        return bool(re.search(r"(.)\1{" + str(threshold - 1) + ",}", text))
+    
+    @staticmethod
+    def has_word_merge(text: str, max_word_len: int = 8) -> bool:
+        words = text.split()
+        if any(len(w) > max_word_len for w in words):
+            return True
+        if len(words) == 1 and len(text) > 6:
+            return True
+        return False
+    
+    @staticmethod
+    def classify(text: str) -> ErrorType:
+        has_rep = ErrorClassifier.has_char_repetition(text)
+        has_merge = ErrorClassifier.has_word_merge(text)
+        has_sub = ErrorClassifier.has_char_substitution(text)
+        error_count = sum([has_rep, has_merge, has_sub])
+        if error_count >= 2:
+            return ErrorType.MIXED
+        elif has_sub:
+            return ErrorType.CHAR_SUBSTITUTION
+        elif has_rep:
+            return ErrorType.CHAR_REPETITION
+        elif has_merge:
+            return ErrorType.WORD_MERGE
+        else:
+            return ErrorType.CLEAN
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# RULES-BASED CORRECTOR
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class RulesBasedCorrector:
+    """Rules-based correction with keyboard proximity mapping."""
+    
+    SUBSTITUTION_MAP = {
+        'ک': 'ك', 'ی': 'ي', 'ے': 'ي',
+        'پ': 'ب', 'چ': 'ج', 'ژ': 'ز',
+        'گ': 'ك', 'ڤ': 'ف', 'ڵ': 'ل',
+        'ڕ': 'ر', 'ڎ': 'د', 'ڼ': 'ن',
+        'ټ': 'ت', 'ډ': 'د', 'ړ': 'ر',
+        'ۀ': 'ه', 'ۃ': 'ة', 'ھ': 'ه',
+        'ە': 'ه', 'ڑ': 'ر'
+    }
+    
+    PREPOSITIONS = {
+        'من', 'في', 'على', 'عن', 'مع', 'إلى', 'الى',
+        'حتى', 'منذ', 'خلال', 'بعد', 'قبل',
+        'ب', 'ل', 'ك', 'لل'
+    }
+    
+    KEYBOARD_NEIGHBORS = {
+        'ض': ['ص', 'ق'], 'ص': ['ض', 'ث', 'ق'], 'ث': ['ص', 'ق'],
+        'ق': ['ض', 'ص', 'ث', 'ف', 'غ'], 'ف': ['ق', 'غ', 'ع', 'ب'],
+        'غ': ['ق', 'ف', 'ع', 'ه'], 'ع': ['ف', 'غ', 'ه', 'خ'],
+        'ه': ['غ', 'ع', 'خ', 'ح'], 'خ': ['ع', 'ه', 'ح', 'ج'],
+        'ح': ['ه', 'خ', 'ج'], 'ج': ['خ', 'ح', 'د'],
+        'د': ['ج', 'ذ'], 'ذ': ['د'],
+        'ش': ['س', 'ي', 'ئ'], 'س': ['ش', 'ي', 'ب'],
+        'ي': ['ش', 'س', 'ب', 'ت'], 'ب': ['ي', 'س', 'ف', 'ل', 'ن'],
+        'ل': ['ب', 'ا', 'ن', 'م'], 'ا': ['ل', 'ت', 'م'],
+        'ت': ['ي', 'ا', 'ن'], 'ن': ['ب', 'ل', 'ت', 'م', 'ك'],
+        'م': ['ل', 'ا', 'ن', 'ك'], 'ك': ['ن', 'م', 'ط'],
+        'ط': ['ك', 'ظ'], 'ظ': ['ط'],
+        'ئ': ['ش', 'ء', 'ر'], 'ء': ['ئ', 'ؤ'], 'ؤ': ['ء', 'ر'],
+        'ر': ['ئ', 'ؤ', 'لا', 'ى', 'ز'], 'لا': ['ر', 'ى'],
+        'ى': ['ر', 'لا', 'ة', 'ز'], 'ة': ['ى', 'و', 'ز'],
+        'و': ['ة', 'ز'], 'ز': ['ر', 'ى', 'ة', 'و'],
+        'أ': ['ا', 'إ', 'آ'], 'إ': ['ا', 'أ'], 'آ': ['ا', 'أ'],
+    }
+    
+    @staticmethod
+    def is_keyboard_neighbor(char1: str, char2: str) -> bool:
+        neighbors = RulesBasedCorrector.KEYBOARD_NEIGHBORS.get(char1, [])
+        return char2 in neighbors
+    
+    @staticmethod
+    def fix_char_substitution(text: str) -> str:
+        for old, new in RulesBasedCorrector.SUBSTITUTION_MAP.items():
+            text = text.replace(old, new)
+        return text
+    
+    @staticmethod
+    def fix_char_repetition(text: str) -> str:
+        text = re.sub(r'([^\d\s])\1{2,}', r'\1', text)
+        return text
+    
+    @staticmethod
+    def advanced_heuristic_repair(text: str) -> str:
+        text = RulesBasedCorrector.fix_char_substitution(text)
+        text = RulesBasedCorrector.fix_char_repetition(text)
+        words = text.split()
+        processed_words = []
+        for word in words:
+            processed_words.append(RulesBasedCorrector._recursive_split(word))
+        return ' '.join(processed_words)
+
+    @staticmethod
+    def _recursive_split(word: str) -> str:
+        if len(word) < 4:
+            return word
+        separables = sorted(['من', 'في', 'على', 'عن', 'مع', 'إلى', 'الى', 'حتى', 'منذ', 'خلال', 'بعد', 'قبل'], key=len, reverse=True)
+        for sep in separables:
+            if word == sep:
+                return word
+            if word.startswith(sep):
+                remainder = word[len(sep):]
+                if len(remainder) >= 3:
+                     return sep + " " + RulesBasedCorrector._recursive_split(remainder)
+        if word.startswith('يا') and len(word) > 4:
+             return 'يا ' + RulesBasedCorrector._recursive_split(word[2:])
+        return word
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# OUTPUT VALIDATOR (Hallucination Prevention)
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class OutputValidator:
+    """Validate model outputs to prevent hallucinations"""
+    
+    @staticmethod
+    def calculate_edit_distance(s1: str, s2: str) -> int:
+        return Levenshtein.distance(s1, s2)
+    
+    @staticmethod
+    def check_character_preservation(original: str, corrected: str) -> Tuple[bool, str]:
+        chars_original = set(original)
+        chars_corrected = set(corrected)
+        if not chars_original:
+            return True, "valid"
+        intersection = chars_original & chars_corrected
+        union = chars_original | chars_corrected
+        jaccard = len(intersection) / len(union) if union else 0
+        if jaccard < 0.35:
+            return False, "low_character_similarity"
+        return True, "valid"
+
+    @staticmethod
+    def check_word_count(original: str, corrected: str) -> Tuple[bool, str]:
+        len_orig = len(original.split())
+        len_corr = len(corrected.split())
+        if len_orig == 1:
+            if len_corr <= 3:
+                return True, "valid"
+            if len(original) > 12 and len_corr <= 6:
+                return True, "valid"
+        ratio = len_corr / len_orig if len_orig > 0 else 0
+        if ratio > 2.0 or ratio < 0.5:
+             return False, "word_count_mismatch"
+        return True, "valid"
+
+    def validate(self, original: str, corrected: str, error_type: str) -> Tuple[bool, str]:
+        if not corrected or not corrected.strip():
+            return False, "empty_output"
+        original_no_space = original.replace(' ', '').replace('\u200c', '')
+        corrected_no_space = corrected.replace(' ', '').replace('\u200c', '')
+        if original_no_space == corrected_no_space:
+            return True, "space_leniency_accept"
+        len_orig = len(original)
+        len_corr = len(corrected)
+        if len_corr > len_orig * 2.5:
+             return False, "too_long"
+        if len_corr < len_orig * 0.5:
+             if error_type == ErrorType.CHAR_REPETITION:
+                 pass
+             else:
+                 return False, "too_short"
+        is_valid_count, reason = self.check_word_count(original, corrected)
+        if not is_valid_count:
+            return False, reason
+        is_valid_chars, reason = self.check_character_preservation(original, corrected)
+        if not is_valid_chars:
+             return False, reason
+        return True, "valid"
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# VOCABULARY MANAGER
+# ══════════════════════════════════════════════════════════════════════════════��
+
+class VocabularyManager:
+    """Centralized vocabulary management for OOV/IV detection."""
+    
+    HAMZA_VARIANTS = {'أ', 'إ', 'آ', 'ء', 'ؤ', 'ئ', 'ا'}
+    ALEF_NORMALIZED = 'ا'
+    TA_MARBUTA = 'ة'
+    HA = 'ه'
+    YA_VARIANTS = {'ي', 'ى'}
+    YA_NORMALIZED = 'ي'
+    
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.vocab = {
+            w for w in tokenizer.get_vocab().keys()
+            if w.isalpha() and not w.startswith('##') and len(w) > 1
+        }
+        self.vocab_rank = {w: i for w, i in tokenizer.get_vocab().items()}
+        self.normalized_vocab = {self.normalize_for_comparison(w): w for w in self.vocab}
+        logger.info(f"VocabularyManager initialized: {len(self.vocab)} words")
+    
+    @classmethod
+    def normalize_for_comparison(cls, word: str) -> str:
+        result = []
+        for i, char in enumerate(word):
+            if char in cls.HAMZA_VARIANTS:
+                result.append(cls.ALEF_NORMALIZED)
+            elif char == cls.TA_MARBUTA and i == len(word) - 1:
+                result.append(cls.HA)
+            elif char in cls.YA_VARIANTS:
+                result.append(cls.YA_NORMALIZED)
+            else:
+                result.append(char)
+        return ''.join(result)
+    
+    def is_iv(self, word: str) -> bool:
+        clean = re.sub(r'[^\w]', '', word)
+        if not clean:
+            return True
+        if clean in self.vocab:
+            return True
+        normalized = self.normalize_for_comparison(clean)
+        if normalized in self.normalized_vocab:
+            return True
+        return False
+    
+    def is_oov(self, word: str) -> bool:
+        return not self.is_iv(word)
+    
+    def get_frequency_rank(self, word: str) -> int:
+        clean = re.sub(r'[^\w]', '', word)
+        return self.vocab_rank.get(clean, 999999)
+    
+    def all_words_iv(self, text: str) -> bool:
+        words = text.split()
+        return all(self.is_iv(w) for w in words)
+    
+    def count_oov_words(self, text: str) -> int:
+        words = text.split()
+        return sum(1 for w in words if self.is_oov(w))
+    
+    def get_oov_words(self, text: str) -> List[str]:
+        words = text.split()
+        return [w for w in words if self.is_oov(w)]
+    
+    def words_are_equivalent(self, word1: str, word2: str) -> bool:
+        norm1 = self.normalize_for_comparison(word1)
+        norm2 = self.normalize_for_comparison(word2)
+        return norm1 == norm2
+    
+    @staticmethod
+    def damerau_levenshtein_distance(s1: str, s2: str) -> int:
+        return jellyfish.damerau_levenshtein_distance(s1, s2)
+    
+    def calculate_similarity(self, original: str, corrected: str) -> float:
+        dist = self.damerau_levenshtein_distance(original, corrected)
+        max_len = max(len(original), len(corrected), 1)
+        return 1.0 - (dist / max_len)
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# WORD ALIGNER
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class WordAligner:
+    """Aligns input and output words to create hybrid corrections."""
+    
+    def __init__(self, vocab_manager):
+        self.vocab = vocab_manager
+    
+    def align_words(self, input_text: str, output_text: str) -> str:
+        input_words = input_text.split()
+        output_words = output_text.split()
+        if abs(len(input_words) - len(output_words)) > 2:
+            input_oov = self.vocab.count_oov_words(input_text)
+            output_oov = self.vocab.count_oov_words(output_text)
+            return output_text if output_oov < input_oov else input_text
+        result = []
+        min_len = min(len(input_words), len(output_words))
+        for i in range(min_len):
+            in_word = input_words[i]
+            out_word = output_words[i]
+            best_word = self._select_best_word(in_word, out_word)
+            result.append(best_word)
+        if len(output_words) > min_len:
+            result.extend(output_words[min_len:])
+        elif len(input_words) > min_len:
+            for w in input_words[min_len:]:
+                 if self.vocab.is_iv(w):
+                     result.append(w)
+        return ' '.join(result)
+    
+    def _select_best_word(self, input_word: str, output_word: str) -> str:
+        if input_word == output_word:
+            return input_word
+        in_iv = self.vocab.is_iv(input_word)
+        out_iv = self.vocab.is_iv(output_word)
+        if not in_iv and out_iv:
+            return output_word
+        if in_iv and not out_iv:
+            return input_word
+        if in_iv and out_iv:
+            return input_word 
+        if len(input_word) == len(output_word) and len(input_word) >= 3:
+            for i in range(len(input_word)):
+                if input_word[i] != output_word[i]:
+                    hybrid = input_word[:i] + output_word[i] + input_word[i+1:]
+                    if self.vocab.is_iv(hybrid):
+                        return hybrid
+                    hybrid2 = output_word[:i] + input_word[i] + output_word[i+1:]
+                    if self.vocab.is_iv(hybrid2):
+                        return hybrid2
+        return output_word
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# SPLIT/MERGE SPECIALIST
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class SplitMergeSpecialist:
+    """Handles word splitting and merging with vocabulary validation."""
+    
+    SEPARABLE_PREFIXES = [
+        'من', 'في', 'على', 'عن', 'مع', 'إلى', 'الى', 'حتى', 'منذ', 'خلال', 
+        'بعد', 'قبل', 'بين', 'حول', 'تحت', 'فوق', 'أمام', 'وراء', 'دون',
+        'أن', 'لن', 'لم', 'قد', 'سوف', 'كي', 'إذا', 'لو', 'مثل', 'غير',
+        'يا',
+    ]
+    
+    PROTECTED_WORDS = {
+        'في', 'من', 'على', 'عن', 'مع', 'إلى', 'الى', 'ان', 'أن', 'لا', 'ما', 'هو', 'هي',
+        'لم', 'لن', 'قد', 'كل', 'كان', 'ذلك', 'هذا', 'هذه', 'التي', 'الذي', 'بين',
+    }
+    
+    ATTACHED_PREFIXES = [
+        'وال', 'بال', 'فال', 'كال', 'لل',
+        'وب', 'وف', 'ول', 'وك', 'وم', 'ون',
+        'فب', 'فل', 'فك', 'فم',
+    ]
+    
+    PRONOUN_SUFFIXES = {'كم', 'هم', 'ها', 'هن', 'كن', 'نا', 'هما', 'كما', 'تم', 'تن'}
+    
+    def __init__(self, vocab_manager):
+        self.vocab = vocab_manager
+        self.separable_prefixes = sorted(
+            self.SEPARABLE_PREFIXES, key=len, reverse=True
+        )
+    
+    def split_word(self, word: str) -> str:
+        if len(word) < 5:
+            return word
+        if self.vocab.is_iv(word):
+            return word
+        if word in self.PROTECTED_WORDS:
+            return word
+        for prefix in self.ATTACHED_PREFIXES:
+            if word.startswith(prefix):
+                remainder = word[len(prefix):]
+                if self.vocab.is_iv(remainder):
+                    return word
+                if prefix.endswith('ال') and self.vocab.is_iv(remainder):
+                    return word
+        for prefix in self.separable_prefixes:
+            if word.startswith(prefix) and len(word) > len(prefix) + 2:
+                remainder = word[len(prefix):]
+                if self.vocab.is_iv(remainder):
+                    return f"{prefix} {remainder}"
+        for i in range(3, len(word) - 2):
+            left = word[:i]
+            right = word[i:]
+            if self.vocab.is_iv(left) and self.vocab.is_iv(right):
+                return f"{left} {right}"
+        return word
+    
+    def merge_fragments(self, text: str) -> str:
+        words = text.split()
+        if len(words) < 2:
+            return text
+        result = []
+        i = 0
+        while i < len(words):
+            word = words[i]
+            if i + 1 < len(words):
+                next_word = words[i + 1]
+                merged = word + next_word
+                if len(next_word) == 1 and next_word in 'ةهاي':
+                    if self.vocab.is_iv(merged):
+                        result.append(merged)
+                        i += 2
+                        continue
+                if word == 'ال' and len(next_word) >= 2:
+                    if self.vocab.is_iv(merged):
+                        result.append(merged)
+                        i += 2
+                        continue
+                if self.vocab.is_oov(word) and self.vocab.is_oov(next_word):
+                    if self.vocab.is_iv(merged):
+                        result.append(merged)
+                        i += 2
+                        continue
+                if len(word) <= 2 and self.vocab.is_oov(word):
+                    if self.vocab.is_iv(merged):
+                        result.append(merged)
+                        i += 2
+                        continue
+                if next_word in self.PRONOUN_SUFFIXES:
+                    if self.vocab.is_iv(merged) and not self.vocab.is_iv(word):
+                        result.append(merged)
+                        i += 2
+                        continue
+                if len(word) <= 3 and len(next_word) <= 3:
+                    if len(merged) >= 5 and self.vocab.is_iv(merged):
+                        result.append(merged)
+                        i += 2
+                        continue
+            result.append(word)
+            i += 1
+        return ' '.join(result)
+    
+    def process_text(self, text: str) -> str:
+        text = self.merge_fragments(text)
+        words = text.split()
+        processed = []
+        for word in words:
+            if self.vocab.is_oov(word) and len(word) >= 4:
+                split_result = self.split_word(word)
+                processed.append(split_result)
+            else:
+                processed.append(word)
+        return ' '.join(processed)
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# EDIT DISTANCE CORRECTOR
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class EditDistanceCorrector:
+    """Generates candidates based on Levenshtein distance."""
+    
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.vocab = {
+            w for w in tokenizer.get_vocab().keys() 
+            if w.isalpha() and not w.startswith('##') and len(w) > 1
+        }
+        self.vocab_rank = {w: i for w, i in tokenizer.get_vocab().items()}
+
+    def edits1(self, word):
+        letters    = 'أابتثجحخدذرزسشصضطظعغفقكلمنهويءآىةئؤ'
+        splits     = [(word[:i], word[i:])    for i in range(len(word) + 1)]
+        deletes    = [L + R[1:]               for L, R in splits if R]
+        transposes = [L + R[1] + R[0] + R[2:] for L, R in splits if len(R)>1]
+        replaces   = [L + c + R[1:]           for L, R in splits if R for c in letters]
+        inserts    = [L + c + R               for L, R in splits for c in letters]
+        return set(deletes + transposes + replaces + inserts)
+
+    def edits2(self, word):
+        return (e2 for e1 in self.edits1(word) for e2 in self.edits1(e1))
+
+    def known(self, words):
+        return set(w for w in words if w in self.vocab)
+
+    def generate_candidate(self, text: str) -> str:
+        words = text.split()
+        corrected_words = []
+        for word in words:
+            clean_word = re.sub(r'[^\w]', '', word)
+            if clean_word in self.vocab:
+                corrected_words.append(word)
+                continue
+            candidates = self.known(self.edits1(clean_word))
+            if not candidates:
+                if len(clean_word) < 7: 
+                    candidates = self.known(self.edits2(clean_word))
+            if candidates:
+                best_candidate = min(candidates, key=lambda w: self.vocab_rank.get(w, 999999))
+                corrected_words.append(best_candidate)
+            else:
+                corrected_words.append(word)
+        return ' '.join(corrected_words)
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# CONTEXTUAL CORRECTOR (MLM-based) — Optional, disabled by default to save RAM
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class ContextualCorrector:
+    """MLM-based contextual correction for confusion pairs"""
+    
+    CONFUSION_PAIRS = [
+        ('ض', 'ظ'), ('ذ', 'ز'), ('ث', 'س'), ('ص', 'س'),
+        ('ط', 'ت'), ('ق', 'ك'), ('ه', 'ة'), ('ا', 'ى'),
+        ('ت', 'د'), ('د', 'ض'), ('ك', 'ق'), ('غ', 'ق'),
+        ('ج', 'ش'), ('س', 'ز'), ('ف', 'ب'), ('و', 'و'),
+        ('ؤ', 'و'), ('ئ', 'ي'), ('ء', 'أ'), ('إ', 'أ'),
+    ]
+    
+    def __init__(self, model_name: str = 'aubmindlab/bert-base-arabertv02', cache_size: int = 10000):
+        from transformers import AutoTokenizer, AutoModelForMaskedLM
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForMaskedLM.from_pretrained(model_name)
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        self.confusion_map = self._build_confusion_map()
+        self.cache_hits = 0
+        self.cache_misses = 0
+        self._score_cache = {}
+        self.cache_size = cache_size
+        self.vocab = self.tokenizer.get_vocab()
+    
+    def _build_confusion_map(self):
+        confusion_map = {}
+        for char1, char2 in self.CONFUSION_PAIRS:
+            if char1 not in confusion_map:
+                confusion_map[char1] = []
+            if char2 not in confusion_map:
+                confusion_map[char2] = []
+            confusion_map[char1].append(char2)
+            confusion_map[char2].append(char1)
+        return confusion_map
+    
+    def get_confusable_chars(self, char: str) -> List[str]:
+        return self.confusion_map.get(char, [])
+    
+    def generate_candidates(self, word: str) -> List[str]:
+        candidates = [word]
+        for i, char in enumerate(word):
+            confusables = self.get_confusable_chars(char)
+            for conf_char in confusables:
+                candidate = word[:i] + conf_char + word[i+1:]
+                if candidate not in candidates:
+                    candidates.append(candidate)
+        for i in range(len(word) - 1):
+            if word[i] == word[i+1]:
+                candidate = word[:i] + word[i+1:]
+                if candidate not in candidates:
+                    candidates.append(candidate)
+        COMMON_CHARS = 'ابتثجحخدذرزسشصضطظعغفقكلمنهويأإآءئؤةى'
+        for i in range(len(word) + 1):
+            for char in COMMON_CHARS:
+                candidate = word[:i] + char + word[i:]
+                if candidate in self.vocab and candidate not in candidates:
+                    candidates.append(candidate)
+        if len(word) < 7:
+            for i in range(len(word)):
+                for char in COMMON_CHARS:
+                    if char != word[i]:
+                        candidate = word[:i] + char + word[i+1:]
+                        if candidate in self.vocab and candidate not in candidates:
+                            candidates.append(candidate)
+        for i in range(len(word)):
+            candidate = word[:i] + word[i+1:]
+            if len(candidate) > 1:
+                if candidate in self.vocab and candidate not in candidates:
+                    candidates.append(candidate)
+        return candidates
+    
+    def score_with_mlm(self, text: str, position: int, word: str) -> float:
+        cache_key = f"{text}|{position}|{word}"
+        if cache_key in self._score_cache:
+            self.cache_hits += 1
+            return self._score_cache[cache_key]
+        self.cache_misses += 1
+        words = text.split()
+        if position >= len(words):
+            return 0.0
+        masked_words = words.copy()
+        masked_words[position] = '[MASK]'
+        masked_text = ' '.join(masked_words)
+        inputs = self.tokenizer(masked_text, return_tensors='pt', padding=True, truncation=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            predictions = outputs.logits
+        mask_token_index = (inputs['input_ids'] == self.tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
+        if len(mask_token_index) == 0:
+            return 0.0
+        mask_token_logits = predictions[0, mask_token_index[0], :]
+        probs = torch.softmax(mask_token_logits, dim=0)
+        word_tokens = self.tokenizer.encode(word, add_special_tokens=False)
+        if not word_tokens:
+            return 0.0
+        word_token_id = word_tokens[0]
+        score = probs[word_token_id].item()
+        if len(self._score_cache) >= self.cache_size:
+            self._score_cache.pop(next(iter(self._score_cache)))
+        self._score_cache[cache_key] = score
+        return score
+    
+    def score_candidates_batch(self, text: str, position: int, candidates: List[str]) -> dict:
+        scores = {}
+        for candidate in candidates:
+            scores[candidate] = self.score_with_mlm(text, position, candidate)
+        return scores
+    
+    def predict_masked_token(self, text: str, position: int, top_k: int = 5) -> List[Tuple[str, float]]:
+        words = text.split()
+        if position >= len(words):
+            return []
+        masked_words = words.copy()
+        masked_words[position] = '[MASK]'
+        masked_text = ' '.join(masked_words)
+        inputs = self.tokenizer(masked_text, return_tensors='pt', padding=True, truncation=True).to(self.device)
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            predictions = outputs.logits
+        mask_token_index = (inputs['input_ids'] == self.tokenizer.mask_token_id).nonzero(as_tuple=True)[1]
+        if len(mask_token_index) == 0:
+            return []
+        mask_token_logits = predictions[0, mask_token_index[0], :]
+        probs = torch.softmax(mask_token_logits, dim=0)
+        top_k_weights, top_k_indices = torch.topk(probs, top_k, sorted=True)
+        results = []
+        for i in range(top_k):
+            token_id = top_k_indices[i].item()
+            score = top_k_weights[i].item()
+            token = self.tokenizer.decode([token_id]).strip()
+            if not token.startswith("##") and token not in self.tokenizer.all_special_tokens:
+                results.append((token, score))
+        return results
+
+    def refine_sentence_with_mask(self, text: str, threshold: float = 0.001, vocab_manager=None, raw_model_output=None) -> str:
+        words = text.split()
+        refined_words = words.copy()
+        raw_words = raw_model_output.split() if raw_model_output else []
+        for i, word in enumerate(words):
+            if vocab_manager and vocab_manager.is_iv(word):
+                continue
+            if i < len(raw_words) and word == raw_words[i]:
+                continue
+            if len(word) <= 2:
+                continue
+            current_score = self.score_with_mlm(text, i, word)
+            if current_score > threshold:
+                continue
+            predictions = self.predict_masked_token(text, i, top_k=10)
+            for pred_word, pred_score in predictions:
+                if pred_word == word:
+                    continue
+                if abs(len(pred_word) - len(word)) > 1:
+                     continue
+                dist = Levenshtein.distance(word, pred_word)
+                max_len = max(len(word), len(pred_word))
+                similarity = 1.0 - (dist / max_len)
+                if similarity < 0.90:
+                    continue
+                if vocab_manager and vocab_manager.is_oov(pred_word):
+                    continue
+                if pred_score < 0.12:
+                    continue
+                is_original_common = current_score > 0.001
+                if is_original_common:
+                     if pred_score > current_score * 1000:
+                         refined_words[i] = pred_word
+                         break
+                else:
+                    if pred_score > current_score * 50 and pred_score > 0.2:
+                        refined_words[i] = pred_word
+                        break
+        return ' '.join(refined_words)
+    
+    def calculate_sentence_score(self, text: str) -> float:
+        words = text.split()
+        if not words:
+            return 0.0
+        total_score = 0.0
+        scored_words = 0
+        for i, word in enumerate(words):
+            score = self.score_with_mlm(text, i, word)
+            total_score += score
+            scored_words += 1
+        if scored_words == 0:
+            return 0.0
+        return total_score / scored_words
+
+
+# ═══════════════════════════════════════════════════════════════════════════════
+# MAIN SPELL CHECKER CLASS
+# ═══════════════════════════════════════════════════════════════════════════════
+
+class ArabicSpellChecker:
+    """Main Arabic Spell Checker class"""
+    
+    def __init__(self, model, tokenizer, device, use_contextual: bool = True):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = device
+        
+        self.postprocessor = AraSpellPostProcessor()
+        self.classifier = ErrorClassifier()
+        self.rules = RulesBasedCorrector()
+        self.validator = OutputValidator()
+        self.vocab_manager = VocabularyManager(tokenizer)
+        self.edit_corrector = EditDistanceCorrector(tokenizer)
+        self.split_merge = SplitMergeSpecialist(self.vocab_manager)
+        self.word_aligner = WordAligner(self.vocab_manager)
+        
+        self.use_contextual = use_contextual
+        if use_contextual:
+            try:
+                self.contextual = ContextualCorrector()
+                logger.info("Contextual correction enabled")
+            except Exception as e:
+                logger.warning(f"Contextual correction disabled: {e}")
+                self.contextual = None
+                self.use_contextual = False
+        else:
+            self.contextual = None
+
+    def _fix_repeated_end_chars(self, text: str) -> str:
+        text = re.sub(r'([ا-ي])\1+\b', r'\1', text)
+        return text
+    
+    def _fix_merged_with_errors(self, text: str) -> str:
+        text = re.sub(r'ال([ا-ي])\1+([ا-ي]{2,})', r'ال\2', text)
+        text = re.sub(r'\b([ا-ي]{3,})([ا-ي])\2+\b', r'\1\2', text)
+        return text
+
+    def _split_merged_words_linguistic(self, text: str) -> str:
+        text = re.sub(
+            r'\b(في|من|إلى|الى|حتى|منذ|خلال|بعد|قبل)(ال)?([ا-ي]{3,})',
+            r'\1 \2\3', text
+        )
+        text = re.sub(r'\b(كل)([ا-ي]{3,})', r'\1 \2', text)
+        text = re.sub(r'([ا-ي]{3,})(ال)([ا-ي]{3,})', r'\1 \2\3', text)
+        text = re.sub(r'\b([بلك])(ال)?([ا-ي]{3,})', r'\1 \2\3', text)
+        text = re.sub(r'([ا-ي]{4,})(عليكم|عليك|عليه|عليها)', r'\1 \2', text)
+        text = re.sub(r'([ا-ي]{3,})(على|عن)([ا-ي]{3,})', r'\1 \2 \3', text)
+        return text
+    
+    def _split_long_words_heuristic(self, text: str, max_length: int = 15) -> str:
+        words = text.split()
+        result = []
+        for word in words:
+            if len(word) <= max_length:
+                result.append(word)
+                continue
+            if 'ال' in word[2:]:
+                parts = word.split('ال', 1)
+                if len(parts[0]) >= 2 and len(parts[1]) >= 3:
+                    result.extend([parts[0], 'ال' + parts[1]])
+                    continue
+            if len(word) >= 8:
+                split_found = False
+                for split_pos in [2, 3]:
+                    prefix = word[:split_pos]
+                    suffix = word[split_pos:]
+                    if prefix in ['في', 'من', 'على', 'عن', 'مع', 'كل', 'ب', 'ل', 'ك']:
+                        result.extend([prefix, suffix])
+                        split_found = True
+                        break
+                if not split_found:
+                    result.append(word)
+            else:
+                result.append(word)
+        return ' '.join(result)
+    
+    def _normalize_tanween_patterns(self, text: str) -> str:
+        text = re.sub(r'([ا-ي]{2,})أ\b', r'\1اً', text)
+        text = re.sub(r'\s+أ\s+', ' ', text)
+        text = re.sub(r'\b([بلك])\s+([ا-ي])', r'\1\2', text)
+        return text
+    
+    def preprocess(self, text: str) -> str:
+        """Preprocessing pipeline"""
+        text = self.postprocessor.remove_harakat(text)
+        text = self.postprocessor.remove_tatweel(text)
+        text = self.postprocessor.normalize_special_chars(text)
+        text = self._fix_repeated_end_chars(text)
+        text = self._fix_merged_with_errors(text)
+        text = self._split_merged_words_linguistic(text)
+        text = self._split_long_words_heuristic(text)
+        text = self._normalize_tanween_patterns(text)
+        text = self.postprocessor.merge_separated_al(text)
+        text = self.postprocessor.unified_collapse_repeated(text)
+        text = self.rules.fix_char_substitution(text)
+        text = self.rules.fix_char_repetition(text)
+        text = self.postprocessor.normalize_spaces(text)
+        return text
+    
+    def postprocess(self, text: str, original: str = "") -> str:
+        """Postprocessing pipeline"""
+        return self.postprocessor.full_postprocess(text, original, vocab_manager=self.vocab_manager)
+    
+    def model_inference(self, text: str, num_return_sequences: int = 5) -> List[str]:
+        """Run seq2seq model inference and return top candidates."""
+        inputs = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=128)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                num_beams=5,
+                num_return_sequences=num_return_sequences,
+                early_stopping=True,
+                return_dict_in_generate=True,
+                output_scores=True
+            )
+        candidates = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)
+        self._last_beam_scores = {}
+        if hasattr(outputs, 'sequences_scores') and outputs.sequences_scores is not None:
+            scores = outputs.sequences_scores.tolist()
+            for cand, score in zip(candidates, scores):
+                self._last_beam_scores[cand] = score
+        return candidates
+    
+    def correct(self, text: str) -> str:
+        """
+        Main correction pipeline (RERANKING APPROACH)
+        
+        Steps:
+        1. Preprocess
+        2. Generate Candidates (Model Beams + Baseline)
+        3. Rerank Candidates (Validator + Fluency)
+        4. Select Best
+        5. Postprocess
+        """
+        if not text or not text.strip():
+            return text
+        
+        original = text
+        
+        # 1. Preprocess
+        preprocessed_text = self.preprocess(text)
+        
+        # 2. Classify error type
+        error_type = self.classifier.classify(preprocessed_text)
+        
+        # 3. Generate Candidates
+        candidates = []
+        candidates.append(preprocessed_text)
+        
+        rules_candidate = self.rules.advanced_heuristic_repair(text)
+        candidates.append(rules_candidate)
+        
+        edit_candidate = self.edit_corrector.generate_candidate(text)
+        if edit_candidate != text and edit_candidate != rules_candidate:
+            candidates.append(edit_candidate)
+        
+        raw_model_output = None
+        try:
+            model_candidates = self.model_inference(preprocessed_text, num_return_sequences=5)
+            raw_model_output = model_candidates[0] if model_candidates else None
+            candidates.extend(model_candidates)
+            
+            if model_candidates:
+                hybrid_candidate = self.word_aligner.align_words(preprocessed_text, model_candidates[0])
+                if hybrid_candidate not in candidates:
+                    candidates.append(hybrid_candidate)
+                for beam in model_candidates[1:3]:
+                    hybrid_beam = self.word_aligner.align_words(preprocessed_text, beam)
+                    if hybrid_beam not in candidates:
+                        candidates.append(hybrid_beam)
+            
+            if model_candidates and len(model_candidates) >= 3:
+                try:
+                    beam_word_lists = [c.split() for c in model_candidates]
+                    max_words = max(len(wl) for wl in beam_word_lists)
+                    voted_words = []
+                    for pos in range(max_words):
+                        words_at_pos = []
+                        for wl in beam_word_lists:
+                            if pos < len(wl):
+                                words_at_pos.append(wl[pos])
+                        if words_at_pos:
+                            most_common = Counter(words_at_pos).most_common(1)[0][0]
+                            voted_words.append(most_common)
+                    voted_candidate = ' '.join(voted_words)
+                    if voted_candidate not in candidates:
+                        candidates.append(voted_candidate)
+                except Exception:
+                    pass
+        except Exception as e:
+            logger.warning(f"Model inference failed: {e}")
+        
+        # Remove duplicates
+        unique_candidates = []
+        seen = set()
+        for c in candidates:
+            if c not in seen:
+                unique_candidates.append(c)
+                seen.add(c)
+        candidates = unique_candidates
+        
+        # 4. Rerank Candidates
+        best_candidate = preprocessed_text
+        best_score = -1.0
+        candidate_scores = []
+        
+        for cand in candidates:
+            is_valid, reason = self.validator.validate(original, cand, error_type.value)
+            if len(cand) < len(original) * 0.5:
+                is_valid = False
+                reason = "too_short"
+
+            input_oov_count = self.vocab_manager.count_oov_words(original)
+            cand_oov_count = self.vocab_manager.count_oov_words(cand)
+            vocab_boost = 1.0
+            
+            if input_oov_count > 0 and cand_oov_count < input_oov_count:
+                oov_reduction = input_oov_count - cand_oov_count
+                vocab_boost = 1.0 + (oov_reduction * 0.3)
+                if cand_oov_count == 0 and self.vocab_manager.all_words_iv(cand):
+                    if not is_valid and reason not in ["empty_output"]:
+                        is_valid = True
+                        reason = "vocab_aware_accept"
+            elif cand_oov_count > input_oov_count:
+                vocab_boost = 0.5
+            elif input_oov_count == 0 and cand_oov_count == 0:
+                vocab_boost = 1.0
+            
+            validity_factor = 1.0 if is_valid else 0.001
+            
+            fluency_score = 0.0
+            if self.use_contextual and self.contextual:
+                try:
+                    fluency_score = self.contextual.calculate_sentence_score(cand)
+                except Exception as e:
+                    logger.warning(f"Scoring failed: {e}")
+                    fluency_score = 0.5
+            else:
+                fluency_score = 1.0 
+            
+            dist = VocabularyManager.damerau_levenshtein_distance(preprocessed_text, cand)
+            max_len = max(len(preprocessed_text), len(cand), 1)
+            similarity = 1.0 - (dist / max_len)
+            if cand == preprocessed_text:
+                similarity = 1.0
+            
+            keyboard_bonus = 1.0
+            input_words = preprocessed_text.split()
+            cand_words = cand.split()
+            if len(input_words) == len(cand_words):
+                for iw, cw in zip(input_words, cand_words):
+                    if iw != cw and len(iw) == len(cw):
+                        for ic, cc in zip(iw, cw):
+                            if ic != cc and RulesBasedCorrector.is_keyboard_neighbor(ic, cc):
+                                keyboard_bonus *= 1.05
+            
+            if fluency_score > 0.85 and cand_oov_count == 0:
+                 if not is_valid and reason in ["too_short", "low_character_similarity", "word_count_mismatch"]:
+                      if len(cand) >= len(original) * 0.4:
+                          is_valid = True
+                          reason = "high_confidence_override"
+                          vocab_boost *= 1.2
+                          validity_factor = 1.0
+            
+            fluency_exp = 0.3
+            similarity_exp = 3.0
+            beam_boost = 1.0
+            if raw_model_output and cand == raw_model_output:
+                beam_boost = 1.15
+            
+            final_score = (fluency_score ** fluency_exp) * (similarity ** similarity_exp) * validity_factor * vocab_boost * keyboard_bonus * beam_boost
+            
+            candidate_scores.append({
+                'text': cand, 'is_valid': is_valid, 'reason': reason,
+                'fluency': fluency_score, 'similarity': similarity,
+                'vocab_boost': vocab_boost, 'input_oov': input_oov_count,
+                'cand_oov': cand_oov_count, 'final_score': final_score
+            })
+            
+            if final_score > best_score:
+                best_score = final_score
+                best_candidate = cand
+        
+        # Output Quality Scoring
+        if best_candidate != preprocessed_text:
+            preprocessed_score = 0.0
+            for cs in candidate_scores:
+                if cs['text'] == preprocessed_text:
+                    preprocessed_score = cs['final_score']
+                    break
+            if preprocessed_score > 0 and best_score < preprocessed_score * 1.05:
+                best_oov = self.vocab_manager.count_oov_words(best_candidate)
+                prep_oov = self.vocab_manager.count_oov_words(preprocessed_text)
+                if best_oov > prep_oov:
+                    best_candidate = preprocessed_text
+                    best_score = preprocessed_score
+        
+        # Contextual Validation Layer
+        if best_candidate != preprocessed_text and self.use_contextual and self.contextual:
+            try:
+                input_fluency = self.contextual.calculate_sentence_score(preprocessed_text)
+                best_fluency = 0.0
+                for cs in candidate_scores:
+                    if cs['text'] == best_candidate:
+                        best_fluency = cs['fluency']
+                        break
+                if input_fluency > 0 and best_fluency > 0:
+                    if input_fluency > best_fluency * 1.5:
+                        input_oov = self.vocab_manager.count_oov_words(preprocessed_text)
+                        best_oov = self.vocab_manager.count_oov_words(best_candidate)
+                        if input_oov <= best_oov:
+                            best_candidate = preprocessed_text
+            except Exception:
+                pass
+        
+        # 5. Postprocess Winner
+        result = self.postprocess(best_candidate, original)
+        
+        # IV-Safe Postprocessing Check
+        if result != best_candidate:
+            result_words = result.split()
+            best_words = best_candidate.split()
+            if len(result_words) == len(best_words):
+                fixed_words = []
+                for idx_fw, (rw, bw) in enumerate(zip(result_words, best_words)):
+                    if rw != bw:
+                        bw_iv = self.vocab_manager.is_iv(bw)
+                        rw_iv = self.vocab_manager.is_iv(rw)
+                        if bw_iv and not rw_iv:
+                            fixed_words.append(bw)
+                        else:
+                            fixed_words.append(rw)
+                    else:
+                        fixed_words.append(rw)
+                result = ' '.join(fixed_words)
+        
+        # 6. Contextual fine-tuning
+        if self.use_contextual and self.contextual:
+             if len(result) > 3:
+                 result = self.contextual.refine_sentence_with_mask(
+                     result, vocab_manager=self.vocab_manager,
+                     raw_model_output=raw_model_output
+                 )
+        
+        # 7. Safe Split/Merge Post-processing
+        result = self.split_merge.merge_fragments(result)
+        
+        # 8. Output Stability Test
+        if result != preprocessed_text and raw_model_output:
+            try:
+                re_preprocessed = self.preprocess(result)
+                stability_dist = VocabularyManager.damerau_levenshtein_distance(result, re_preprocessed)
+                result_len = max(len(result), 1)
+                if stability_dist > 0:
+                    stability_ratio = stability_dist / result_len
+                    if stability_ratio > 0.15:
+                        raw_re = self.preprocess(raw_model_output)
+                        raw_stability = VocabularyManager.damerau_levenshtein_distance(
+                            raw_model_output, raw_re
+                        ) / max(len(raw_model_output), 1)
+                        if raw_stability < stability_ratio:
+                            raw_oov = self.vocab_manager.count_oov_words(raw_model_output)
+                            our_oov = self.vocab_manager.count_oov_words(result)
+                            if raw_oov <= our_oov:
+                                result = raw_model_output
+            except Exception:
+                pass
+        
+        # 9. Bidirectional Word-Level Validation
+        if raw_model_output and result != raw_model_output:
+            result_words = result.split()
+            raw_words = raw_model_output.split()
+            if len(result_words) == len(raw_words):
+                corrected_words = []
+                changed = False
+                for rw, raw_w in zip(result_words, raw_words):
+                    if rw != raw_w:
+                        rw_iv = self.vocab_manager.is_iv(rw)
+                        raw_iv = self.vocab_manager.is_iv(raw_w)
+                        if not rw_iv and raw_iv:
+                            corrected_words.append(raw_w)
+                            changed = True
+                        elif rw_iv and raw_iv:
+                            input_words_list = preprocessed_text.split()
+                            idx = len(corrected_words)
+                            if idx < len(input_words_list):
+                                input_w = input_words_list[idx]
+                                rw_dist = Levenshtein.distance(input_w, rw)
+                                raw_dist = Levenshtein.distance(input_w, raw_w)
+                                if raw_dist < rw_dist:
+                                    corrected_words.append(raw_w)
+                                    changed = True
+                                else:
+                                    corrected_words.append(rw)
+                            else:
+                                corrected_words.append(rw)
+                        else:
+                            corrected_words.append(rw)
+                    else:
+                        corrected_words.append(rw)
+                if changed:
+                    new_result = ' '.join(corrected_words)
+                    new_oov = self.vocab_manager.count_oov_words(new_result)
+                    old_oov = self.vocab_manager.count_oov_words(result)
+                    if new_oov <= old_oov:
+                        result = new_result
+        
+        # 10. SAFETY NET
+        if raw_model_output and raw_model_output != result:
+            raw_oov = self.vocab_manager.count_oov_words(raw_model_output)
+            our_oov = self.vocab_manager.count_oov_words(result)
+            if raw_oov == 0 and our_oov > 0:
+                is_valid, reason = self.validator.validate(original, raw_model_output, "mixed")
+                if is_valid or reason == "space_leniency_accept":
+                    result = raw_model_output
+            elif raw_oov == 0 and our_oov == 0:
+                raw_dist = VocabularyManager.damerau_levenshtein_distance(original, raw_model_output)
+                our_dist = VocabularyManager.damerau_levenshtein_distance(original, result)
+                result_vs_raw_dist = VocabularyManager.damerau_levenshtein_distance(result, raw_model_output)
+                if raw_dist < our_dist and result_vs_raw_dist <= 3:
+                    raw_valid, _ = self.validator.validate(original, raw_model_output, "mixed")
+                    if raw_valid:
+                        result = raw_model_output
+            elif raw_oov == 0:
+                raw_wc = len(raw_model_output.split())
+                our_wc = len(result.split())
+                if raw_wc != our_wc:
+                    raw_dist = VocabularyManager.damerau_levenshtein_distance(original, raw_model_output)
+                    our_dist = VocabularyManager.damerau_levenshtein_distance(original, result)
+                    if raw_dist < our_dist:
+                        raw_valid, _ = self.validator.validate(original, raw_model_output, "mixed")
+                        if raw_valid:
+                            result = raw_model_output
+        
+        return result

src/nlp/spelling/araspell_service.py

@@ -0,0 +1,105 @@
+"""
+AraSpell Service — Lazy-loaded Arabic spelling correction.
+
+Model is loaded on first request and kept in memory.
+Pre-downloaded during Docker build; loaded from HF cache at runtime (no network needed).
+"""
+
+import os
+import logging
+import time
+import torch
+
+logger = logging.getLogger(__name__)
+
+# ── Lazy-loaded singletons ──
+_spell_checker = None
+_load_error = None
+
+# Model identifiers
+MODEL_REPO = 'bayan10/AraSpell-Model'
+MODEL_FILENAME = 'last_model.pt'
+TOKENIZER_NAME = 'aubmindlab/bert-base-arabertv02'
+
+
+def get_spelling_model():
+    """
+    Lazy-load the spelling model on first call.
+    Returns the ArabicSpellChecker instance, or raises RuntimeError if loading fails.
+    """
+    global _spell_checker, _load_error
+
+    if _spell_checker is not None:
+        return _spell_checker
+
+    if _load_error is not None:
+        raise RuntimeError(f"Spelling model previously failed to load: {_load_error}")
+
+    try:
+        t0 = time.time()
+        logger.info("Loading AraSpell spelling model (lazy init)...")
+
+        from huggingface_hub import hf_hub_download
+        from transformers import AutoTokenizer, EncoderDecoderModel
+
+        # 1. Download checkpoint (from HF cache — pre-downloaded in Docker build)
+        logger.info(f"Resolving checkpoint: {MODEL_REPO}/{MODEL_FILENAME}")
+        model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
+        logger.info(f"Checkpoint path: {model_path}")
+
+        # 2. Load tokenizer
+        logger.info(f"Loading tokenizer: {TOKENIZER_NAME}")
+        tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+
+        # 3. Build encoder-decoder model from AraBERT
+        logger.info("Building EncoderDecoderModel from AraBERT...")
+        model = EncoderDecoderModel.from_encoder_decoder_pretrained(
+            TOKENIZER_NAME, TOKENIZER_NAME
+        )
+
+        # 4. Configure generation
+        model.config.decoder_start_token_id = tokenizer.cls_token_id
+        model.config.pad_token_id = tokenizer.pad_token_id
+        model.config.eos_token_id = tokenizer.sep_token_id
+        model.generation_config.max_length = 128
+        model.generation_config.decoder_start_token_id = tokenizer.cls_token_id
+        model.generation_config.pad_token_id = tokenizer.pad_token_id
+        model.generation_config.eos_token_id = tokenizer.sep_token_id
+
+        # 5. Load trained weights
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        logger.info(f"Loading checkpoint weights on {device}...")
+        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+        model.load_state_dict(checkpoint['model_state_dict'], strict=False)
+        model = model.to(device)
+        model.eval()
+
+        epoch = checkpoint.get('epoch', 'N/A')
+        logger.info(f"Spelling model loaded on {device}, epoch: {epoch}")
+
+        # 6. Initialize the spell checker pipeline (contextual=False to save RAM)
+        from nlp.spelling.araspell_rules import ArabicSpellChecker
+        _spell_checker = ArabicSpellChecker(
+            model, tokenizer, device, use_contextual=False
+        )
+
+        elapsed = time.time() - t0
+        logger.info(f"AraSpell ready in {elapsed:.1f}s")
+        return _spell_checker
+
+    except Exception as e:
+        import traceback
+        _load_error = str(e)
+        logger.error(f"Failed to load spelling model: {e}")
+        logger.error(traceback.format_exc())
+        raise RuntimeError(f"Spelling model load failed: {e}")
+
+
+def is_loaded() -> bool:
+    """Check if the spelling model is loaded."""
+    return _spell_checker is not None
+
+
+def get_load_error() -> str:
+    """Return the last load error, or empty string."""
+    return _load_error or ""