""" AutoComplete Service — Hybrid bigram + GPT-2 Arabic autocomplete. COMPLETELY INDEPENDENT from the correction pipeline. This module has ZERO interaction with: - /api/analyze - StageLockManager / OffsetMapper / ClaimedRanges - OverlapResolver / PatchSet / CorrectionPatch - Highlight rendering Architecture: User types → debounce → POST /api/autocomplete → HybridAutoComplete.predict(context) → Bigram lookup + GPT-2 scoring → Ranked suggestions returned to frontend """ import os import time import pickle import logging import threading from functools import lru_cache import torch from huggingface_hub import hf_hub_download logger = logging.getLogger(__name__) # ─── Singleton ──────────────────────────────────────────────────────────────── _instance = None _lock = threading.Lock() def get_autocomplete_model(): """Lazy-loaded singleton — returns the cached HybridAutoComplete instance.""" global _instance if _instance is not None: return _instance with _lock: if _instance is not None: return _instance _instance = HybridAutoComplete() return _instance def is_loaded() -> bool: """Check if the autocomplete model is loaded (without triggering lazy load).""" return _instance is not None and _instance.is_ready() # ─── Cache key helper ───────────────────────────────────────────────────────── def _context_key(context: str) -> str: """Use last 5 words for cache key — preserves enough context for GPT-2 awareness.""" words = context.strip().split() return " ".join(words[-5:]) if words else "" # ─── Main Service ───────────────────────────────────────────────────────────── class HybridAutoComplete: """ Hybrid Arabic autocomplete: 1. Statistical (bigram) — fast, always available 2. Neural (GPT-2) — contextual, optional (may OOM on free tier) 3. Hybrid scoring: alpha * stat + (1-alpha) * neural """ BIGRAM_REPO = "bayan10/AutoComplete" BIGRAM_FILE = "bigram_model_v4.pkl" GPT2_MODEL = "aubmindlab/aragpt2-base" def __init__(self): t0 = time.time() logger.info("Loading AutoComplete model (lazy init)...") self.unigrams = None self.bigrams = None self.gpt2_tokenizer = None self.gpt2_model = None self.device = "cpu" self.alpha = 0.4 # Weight: 40% bigram, 60% GPT-2 (GPT-2 has context!) self.threshold = 0.05 # Min score to show a suggestion self._cache = {} self._cache_max = 256 # 1. Load bigram (required — small file) self._load_bigram() # 2. Load GPT-2 (optional — large model, may OOM) self._load_gpt2() elapsed = time.time() - t0 mode = "hybrid" if self.gpt2_model else "bigram-only" logger.info(f"AutoComplete ready in {elapsed:.1f}s (mode: {mode})") def _load_bigram(self): """Load bigram model from HuggingFace Hub.""" try: path = hf_hub_download( repo_id=self.BIGRAM_REPO, filename=self.BIGRAM_FILE, ) with open(path, "rb") as f: data = pickle.load(f) self.unigrams = data["unigrams"] self.bigrams = data["bigrams"] self._top_unigrams = sorted( [(w, c) for w, c in self.unigrams.items() if len(w) >= 2], key=lambda x: x[1], reverse=True )[:200] logger.info( f"Bigram model loaded: {len(self.unigrams)} unigrams, " f"{len(self.bigrams)} bigram contexts" ) except Exception as e: logger.error(f"Failed to load bigram model: {e}") self.unigrams = {} self.bigrams = {} self._top_unigrams = [] def _load_gpt2(self): """Load GPT-2 model with OOM fallback.""" try: from transformers import GPT2LMHeadModel, GPT2Tokenizer logger.info(f"Loading GPT-2 tokenizer: {self.GPT2_MODEL}") self.gpt2_tokenizer = GPT2Tokenizer.from_pretrained(self.GPT2_MODEL) self.gpt2_tokenizer.pad_token = self.gpt2_tokenizer.eos_token logger.info(f"Loading GPT-2 model: {self.GPT2_MODEL}") self.gpt2_model = GPT2LMHeadModel.from_pretrained(self.GPT2_MODEL) self.gpt2_model.config.pad_token_id = self.gpt2_tokenizer.eos_token_id self.gpt2_model.eval() logger.info("GPT-2 loaded successfully (hybrid mode enabled)") except (torch.cuda.OutOfMemoryError, MemoryError, RuntimeError) as e: logger.warning(f"GPT-2 OOM — falling back to bigram-only mode: {e}") self.gpt2_tokenizer = None self.gpt2_model = None except Exception as e: logger.warning(f"GPT-2 load failed — bigram-only mode: {e}") self.gpt2_tokenizer = None self.gpt2_model = None # ─── Prediction ─────────────────────────────────────────────────────── def predict(self, context: str, n: int = 3) -> list: """ Get top-N autocomplete suggestions for the given context. Args: context: Text before the cursor (last ~200 chars) n: Number of suggestions to return Returns: List of suggestion strings (ranked by score) """ if not context or not context.strip(): return [] context = context.strip() # Check cache cache_key = _context_key(context) if cache_key in self._cache: return self._cache[cache_key][:n] try: if self.gpt2_model is not None: results = self._hybrid_predict(context, n) else: results = self._bigram_predict(context, n) # Cache the result if len(self._cache) >= self._cache_max: # Evict oldest entries (simple FIFO) keys = list(self._cache.keys()) for k in keys[:len(keys) // 2]: del self._cache[k] self._cache[cache_key] = results return results[:n] except Exception as e: logger.error(f"AutoComplete prediction error: {e}") return [] def _bigram_predict(self, context: str, n: int = 3) -> list: """Statistical-only prediction using bigram model.""" from .autocomplete_rules import merge_similar_predictions, filter_suggestions tokens = context.strip().split() if not tokens: return [] last_word = tokens[-1] candidates = [] # Try bigram lookup if last_word in self.bigrams: for w, c in self.bigrams[last_word].items(): if len(w) < 2 or w == last_word: continue candidates.append((w, c)) # Fallback to precomputed top unigrams if no bigram matches if not candidates: candidates = list(self._top_unigrams) if not candidates: return [] total = sum(c for _, c in candidates) if total == 0: return [] preds = [(w, c / total) for w, c in candidates] preds.sort(key=lambda x: x[1], reverse=True) preds = merge_similar_predictions(preds, top_k=n * 3) preds = filter_suggestions(preds) # Apply score threshold preds = [(w, s) for w, s in preds if s >= self.threshold] return [w for w, _ in preds[:n]] def _hybrid_predict(self, context: str, n: int = 3) -> list: """Hybrid prediction: bigram + GPT-2 scoring. GPT-2 receives the FULL sentence as context for true context awareness. Bigram provides frequency-based candidates from the last word. GPT-2's own top predictions are ADDED as candidates so contextually appropriate words that bigram doesn't know about can still appear. """ from .autocomplete_rules import merge_similar_predictions, filter_suggestions tokens = context.strip().split() if not tokens: return [] last_word = tokens[-1] # 1. Get bigram candidates (frequency-based, last word only) stat_candidates = [] if last_word in self.bigrams: for w, c in self.bigrams[last_word].items(): if len(w) < 2 or w == last_word: continue stat_candidates.append((w, c)) # 2. Get GPT-2 next-token probabilities using FULL context # GPT-2 sees the entire sentence, not just the last word gpt2_probs = self._gpt2_next_token_probs(context, top_k=50) # 3. If no bigram candidates, use GPT-2 predictions directly if not stat_candidates: if gpt2_probs: # Use GPT-2's own contextual predictions gpt2_preds = sorted(gpt2_probs.items(), key=lambda x: x[1], reverse=True) gpt2_preds = [(w, s) for w, s in gpt2_preds if s >= self.threshold] gpt2_preds = filter_suggestions(gpt2_preds) return [w for w, _ in gpt2_preds[:n]] return self._bigram_predict(context, n) total = sum(c for _, c in stat_candidates) if total == 0: if gpt2_probs: gpt2_preds = sorted(gpt2_probs.items(), key=lambda x: x[1], reverse=True) gpt2_preds = [(w, s) for w, s in gpt2_preds if s >= self.threshold] gpt2_preds = filter_suggestions(gpt2_preds) return [w for w, _ in gpt2_preds[:n]] return self._bigram_predict(context, n) stat_preds = [(w, c / total) for w, c in stat_candidates] stat_preds.sort(key=lambda x: x[1], reverse=True) stat_preds = merge_similar_predictions(stat_preds, top_k=20) # 4. Hybrid scoring: combine bigram frequency with GPT-2 context score results = [] seen_words = set() for w, stat_p in stat_preds: neural_p = gpt2_probs.get(w, 1e-8) score = self.alpha * stat_p + (1 - self.alpha) * neural_p results.append((w, score)) seen_words.add(w) # 5. ADD GPT-2's own top contextual predictions as bonus candidates # These are complete words GPT-2 generated based on full sentence context for w, neural_p in sorted(gpt2_probs.items(), key=lambda x: x[1], reverse=True)[:10]: if w not in seen_words and len(w) >= 2: score = 0.5 * neural_p # High weight — GPT-2 understands context results.append((w, score)) seen_words.add(w) results.sort(key=lambda x: x[1], reverse=True) results = filter_suggestions(results) # Apply score threshold — filter out low-confidence suggestions results = [(w, s) for w, s in results if s >= self.threshold] return [w for w, _ in results[:n]] def _gpt2_next_token_probs(self, prefix: str, top_k: int = 50) -> dict: """ Get GPT-2 next-WORD predictions using generate() for complete words. Uses SAMPLING (not beam search) for diverse predictions. Beam search collapses: 10 beams → same 1-3 words. Sampling with top_k/top_p → 10-15 diverse, contextual words. """ if self.gpt2_model is None or self.gpt2_tokenizer is None: return {} import re ARABIC_WORD_RE = re.compile(r'[\u0600-\u06FF]{2,}') try: inputs = self.gpt2_tokenizer( prefix, return_tensors="pt", truncation=True, max_length=512, ) input_len = inputs['input_ids'].shape[1] # Generate diverse sequences using SAMPLING (not beam search) with torch.no_grad(): outputs = self.gpt2_model.generate( **inputs, max_new_tokens=5, # Enough for 1 complete Arabic word do_sample=True, # Sampling for diversity top_k=50, # Consider top 50 tokens top_p=0.9, # Nucleus sampling temperature=0.8, # Slight sharpening for quality num_return_sequences=15, # More samples = more diverse words no_repeat_ngram_size=2, ) # Extract the first NEW Arabic word from each sequence word_counts = {} for seq in outputs: new_tokens = seq[input_len:] generated_text = self.gpt2_tokenizer.decode(new_tokens, skip_special_tokens=True).strip() if not generated_text: continue match = ARABIC_WORD_RE.search(generated_text) if match: word = match.group(0) word_counts[word] = word_counts.get(word, 0) + 1 # Score = frequency across samples (how often GPT-2 picks this word) total = sum(word_counts.values()) if total == 0: return {} prob_dict = {w: count / total for w, count in word_counts.items()} logger.info(f"[GPT2] context='{prefix[-50:]}' → words={sorted(prob_dict.items(), key=lambda x: -x[1])[:8]}") return prob_dict except Exception as e: logger.warning(f"GPT-2 scoring failed: {e}") return {} # ─── Health ─────────────────────────────────────────────────────────── def is_ready(self) -> bool: """Returns True if at least the bigram model is loaded.""" return bool(self.unigrams) def get_mode(self) -> str: """Returns 'hybrid', 'bigram-only', or 'unavailable'.""" if self.gpt2_model and self.unigrams: return "hybrid" elif self.unigrams: return "bigram-only" return "unavailable"