""" Punctuation Service — Lazy-loaded Arabic punctuation restoration. Uses: 1. bayan10/PuncAra-v1 (EncoderDecoderModel — local, seq2seq) 2. Rule-based pre/post-processing from punctuation_rules.py Model loaded on first request and kept in memory. """ import logging import time import torch import re import threading logger = logging.getLogger(__name__) # ── Lazy-loaded singletons ── _punctuation_checker = None _load_error = None _lock = threading.Lock() HF_REPO_ID = "bayan10/PuncAra-v1" class PunctuationChecker: """ Arabic punctuation restoration pipeline: 1. Preprocessing (remove diacritics) 2. Model inference (chunked, windowed — 50 words/chunk) 3. Postprocessing: strip non-punctuation changes (Fix P1) 4. Typographic cleanup """ # Arabic and common punctuation marks PUNCTUATION_CHARS = set('.,;:!?،؛؟!.:«»"\'()-–—…') def __init__(self, model, tokenizer, device): self.model = model self.tokenizer = tokenizer self.device = device @staticmethod def _strip_punct(word: str) -> str: """Remove leading/trailing punctuation from a word.""" return word.strip('.,;:!?،؛؟!.:«»"\'()-–—…') @staticmethod def _normalize_hamza(word: str) -> str: import re return re.sub(r'[أإآ]', 'ا', word).replace('ة', 'ه').replace('ى', 'ي') def _strip_non_punctuation_changes(self, original: str, punctuated: str) -> str: """ Fix P1: The PuncAra model was fine-tuned on data with spelling/grammar corrections. We only want punctuation marks from this stage. Strategy: Align original and punctuated word-by-word. For each word, if the model changed the BASE text (not just added/moved punctuation), revert to the original word but keep any punctuation the model added. """ # CRITICAL FIX: The model sometimes predicts standalone punctuation tokens (e.g. "المعلم :"). # If we split by space, the arrays will misalign. We MUST attach punctuation to the preceding word # before splitting to ensure 1-to-1 alignment. import re punctuated_normalized = re.sub(r'\s+([،؛:!؟.])', r'\1', punctuated) orig_words = original.split() punc_words = punctuated_normalized.split() if not orig_words or not punc_words: return punctuated # Build result by aligning words result = [] oi = 0 # index into orig_words pi = 0 # index into punc_words while oi < len(orig_words) and pi < len(punc_words): o_word = orig_words[oi] p_word = punc_words[pi] o_base = self._strip_punct(o_word) p_base = self._strip_punct(p_word) if o_base == p_base or self._normalize_hamza(o_base) == self._normalize_hamza(p_base): # Anti-hallucination for question marks if '؟' in p_word and '؟' not in o_word: _EXCL_CUES = {'هل', 'أين', 'متى', 'كيف', 'لماذا', 'ماذا', 'أي', 'كم', 'ما'} if not any(w in _EXCL_CUES for w in orig_words): p_word = p_word.replace('؟', '.') # Same base word — keep punctuation changes from model result.append(p_word) oi += 1 pi += 1 elif self._is_only_punct_difference(o_word, p_word): if '؟' in p_word and '؟' not in o_word: _EXCL_CUES = {'هل', 'أين', 'متى', 'كيف', 'لماذا', 'ماذا', 'أي', 'كم', 'ما'} if not any(w in _EXCL_CUES for w in orig_words): p_word = p_word.replace('؟', '.') # Words differ only by punctuation — keep model's punctuation result.append(p_word) oi += 1 pi += 1 else: # Model changed the actual word content (spelling/grammar/hamza) # Revert to original word but transfer any NEW punctuation punct_suffix = '' punct_prefix = '' for ch in reversed(p_word): if ch in self.PUNCTUATION_CHARS: punct_suffix = ch + punct_suffix else: break for ch in p_word: if ch in self.PUNCTUATION_CHARS: punct_prefix += ch else: break # Only add punctuation that wasn't already there if not o_word.endswith(punct_suffix) and punct_suffix: if '؟' in punct_suffix and '؟' not in o_word: _EXCL_CUES = {'هل', 'أين', 'متى', 'كيف', 'لماذا', 'ماذا', 'أي', 'كم', 'ما'} if not any(w in _EXCL_CUES for w in orig_words): punct_suffix = punct_suffix.replace('؟', '.') result.append(o_word + punct_suffix) elif punct_prefix and not o_word.startswith(punct_prefix): result.append(punct_prefix + o_word) else: result.append(o_word) oi += 1 pi += 1 # Append remaining original words while oi < len(orig_words): result.append(orig_words[oi]) oi += 1 # Append remaining punctuation-only words from model while pi < len(punc_words): p_word = punc_words[pi] if all(ch in self.PUNCTUATION_CHARS or ch.isspace() for ch in p_word): if '؟' in p_word: _EXCL_CUES = {'هل', 'أين', 'متى', 'كيف', 'لماذا', 'ماذا', 'أي', 'كم', 'ما'} if not any(w in _EXCL_CUES for w in orig_words): p_word = p_word.replace('؟', '.') result.append(p_word) pi += 1 return ' '.join(result) @staticmethod def _is_only_punct_difference(word1: str, word2: str) -> bool: """Check if two words differ only by punctuation characters.""" PUNCT = set('.,;:!?،؛؟!.:«»"\'()-–—…') base1 = ''.join(c for c in word1 if c not in PUNCT) base2 = ''.join(c for c in word2 if c not in PUNCT) return base1 == base2 def _predict_chunk(self, text_chunk: str) -> str: """Run model inference on a single chunk (max 128 tokens).""" from nlp.punctuation.punctuation_rules import arabic_preprocessing text_chunk = arabic_preprocessing(text_chunk) inputs = self.tokenizer( text_chunk, return_tensors="pt", padding=True, truncation=True, max_length=128 ).to(self.device) with torch.no_grad(): outputs = self.model.generate( inputs.input_ids, attention_mask=inputs.attention_mask, decoder_start_token_id=self.tokenizer.cls_token_id, bos_token_id=self.tokenizer.cls_token_id, eos_token_id=self.tokenizer.sep_token_id, pad_token_id=self.tokenizer.pad_token_id, max_length=128, num_beams=3, repetition_penalty=1.2, length_penalty=1.0, early_stopping=True, do_sample=False ) return self.tokenizer.decode(outputs[0], skip_special_tokens=True) def _fix_punctuation(self, text: str) -> str: """Process a paragraph using non-overlapping window chunking.""" words = text.split() total_words = len(words) window_size = 50 stride = 50 if total_words <= window_size: return self._predict_chunk(text) segments_output = [] for i in range(0, total_words, stride): chunk_words = words[i: i + window_size] chunk_text = " ".join(chunk_words) if not chunk_text.strip(): continue processed_segment = self._predict_chunk(chunk_text).strip() # Remove trailing punctuation from non-last segments (context continues) is_last_segment = (i + window_size) >= total_words if not is_last_segment: punctuation_marks = ".?!،؛:؟!" if processed_segment and processed_segment[-1] in punctuation_marks: processed_segment = processed_segment[:-1] segments_output.append(processed_segment) result = " ".join(segments_output) result = re.sub(r'\s+', ' ', result).strip() return result def correct(self, text: str) -> str: """ Run full punctuation restoration on text. Handles multi-paragraph documents. Returns punctuated text, or original text on failure. """ if not text or not text.strip(): return text try: from nlp.punctuation.punctuation_rules import arabic_postprocessing # Split into paragraphs paragraphs = [p.strip() for p in text.split('\n') if p.strip()] processed_paragraphs = [] for paragraph in paragraphs: punctuated = self._fix_punctuation(paragraph) # Fix P1: Strip spelling/grammar changes, keep only punctuation punctuated = self._strip_non_punctuation_changes(paragraph, punctuated) cleaned = arabic_postprocessing(punctuated) processed_paragraphs.append(cleaned) result = "\n".join(processed_paragraphs) _r_display = result[:80] + ('...' if len(result) > 80 else '') _t_display = text[:80] + ('...' if len(text) > 80 else '') logger.info(f"Punctuation output: '{_r_display}' (input: '{_t_display}')") return result except Exception as e: logger.error(f"Punctuation correction failed: {e}") return text def get_punctuation_model(): """ Lazy-load the punctuation model on first call. Returns the PunctuationChecker instance, or raises RuntimeError if loading fails. """ global _punctuation_checker, _load_error if _punctuation_checker is not None: return _punctuation_checker with _lock: if _punctuation_checker is not None: return _punctuation_checker if _load_error is not None: raise RuntimeError(f"Punctuation model previously failed to load: {_load_error}") try: t0 = time.time() logger.info("Loading PuncAra-v1 punctuation model (lazy init)...") device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logger.info(f"Punctuation model device: {device}") from transformers import EncoderDecoderModel, AutoTokenizer logger.info(f"Loading model from HF Hub: {HF_REPO_ID}") model = EncoderDecoderModel.from_pretrained(HF_REPO_ID) tokenizer = AutoTokenizer.from_pretrained(HF_REPO_ID) # Configure special tokens model.config.decoder_start_token_id = tokenizer.cls_token_id model.config.bos_token_id = tokenizer.cls_token_id model.config.eos_token_id = tokenizer.sep_token_id model.config.pad_token_id = tokenizer.pad_token_id model = model.to(device) model.eval() _punctuation_checker = PunctuationChecker(model, tokenizer, device) elapsed = time.time() - t0 logger.info(f"PuncAra-v1 ready in {elapsed:.1f}s") return _punctuation_checker except Exception as e: import traceback _load_error = str(e) logger.error(f"Failed to load punctuation model: {e}") logger.error(traceback.format_exc()) raise RuntimeError(f"Punctuation model load failed: {e}") def is_loaded() -> bool: """Check if the punctuation model is loaded.""" return _punctuation_checker is not None def get_load_error() -> str: """Return the last load error, or empty string.""" return _load_error or ""