vazish/query_norm
0
1"""2Query Normalization Benchmark3==============================4Benchmarks multiple normalization approaches on the generated dataset.5 6Normalizers:7 1. Identity - baseline, no change8 2. PySpellChecker - token-by-token spell correction (current approach)9 3. SymSpell - faster, supports compound word correction10 4. Rules - regex + entity canonicalization (flight IDs, stock tickers, product spacing)11 5. RapidFuzz - fuzzy brand name matching12 6. Combined - Rules → SymSpell → RapidFuzz pipeline13 --- ML ---14 7. ContextualSpellCheck - spaCy pipeline with BERT contextual embeddings15 8. T5SpellCorrector - HuggingFace T5 fine-tuned for spelling correction16 9. CombinedML - Rules → T5 pipeline (entity rules first, T5 for the rest)17 18Metrics (per normalizer, per category):19 exact_match - % where output == canonical (case-insensitive)20 cer - character error rate: edit_dist / max(len_pred, len_gold)21 wer - word error rate: token-level edit distance / n_gold_tokens22 no_change_precision - on no_change rows: % correctly left unchanged23 over_correction - on no_change rows: % wrongly changed24 latency_mean_ms - mean per-query latency25 latency_p50_ms - p50 latency26 latency_p95_ms - p95 latency27 latency_p99_ms - p99 latency28 29Usage:30 pip install -r requirements.txt31 python3 benchmark.py [--dataset dataset.csv]32"""33 34import re35import sys36import time37import argparse38import warnings39import numpy as np40import pandas as pd41from pathlib import Path42from abc import ABC, abstractmethod43from typing import Optional44 45warnings.filterwarnings("ignore")46 47# ── Optional imports ───────────────────────────────────────────────────────────48 49try:50 from Levenshtein import distance as _lev51 def edit_distance(a: str, b: str) -> int: return _lev(a, b)52except ImportError:53 # Pure-python fallback54 def edit_distance(a: str, b: str) -> int:55 m, n = len(a), len(b)56 dp = list(range(n + 1))57 for i in range(1, m + 1):58 prev = dp[:]59 dp[0] = i60 for j in range(1, n + 1):61 dp[j] = prev[j - 1] if a[i-1] == b[j-1] else 1 + min(prev[j], dp[j-1], prev[j-1])62 return dp[n]63 64try:65 from spellchecker import SpellChecker as _SC66 HAS_PYSPELL = True67except ImportError:68 HAS_PYSPELL = False69 print("Warning: pyspellchecker not installed — skipping PySpell normalizer")70 71try:72 from symspellpy import SymSpell as _SS, Verbosity as _V73 import pkg_resources74 HAS_SYMSPELL = True75except ImportError:76 HAS_SYMSPELL = False77 print("Warning: symspellpy not installed — skipping SymSpell normalizer")78 79try:80 from rapidfuzz import process as _rf_process, fuzz as _rf_fuzz81 HAS_RAPIDFUZZ = True82except ImportError:83 HAS_RAPIDFUZZ = False84 print("Warning: rapidfuzz not installed — skipping RapidFuzz normalizer")85 86try:87 import spacy as _spacy88 import contextualSpellCheck as _csc89 _csc_nlp = _spacy.load("en_core_web_sm")90 _csc.add_to_pipe(_csc_nlp)91 HAS_CONTEXTUAL = True92except Exception:93 HAS_CONTEXTUAL = False94 print("Warning: contextualSpellCheck/spacy not available — skipping ContextualSpellCheck normalizer")95 print(" Install: pip install contextualSpellCheck && python -m spacy download en_core_web_sm")96 97try:98 from transformers import pipeline as _hf_pipeline99 HAS_TRANSFORMERS = True100except ImportError:101 HAS_TRANSFORMERS = False102 print("Warning: transformers not installed — skipping T5 normalizer")103 print(" Install: pip install transformers torch")104 105# ── Brand list for fuzzy matching ──────────────────────────────────────────────106 107BRANDS = [108 "amazon", "google", "facebook", "twitter", "instagram", "youtube",109 "linkedin", "reddit", "netflix", "spotify", "microsoft", "adobe",110 "dropbox", "github", "slack", "zoom", "paypal", "ebay", "walmart",111 "target", "best buy", "new york times", "bbc", "cnn", "espn",112 "gmail", "outlook", "yahoo", "apple", "samsung", "dell", "hp",113 "lenovo", "asus", "acer", "toshiba", "sony", "lg", "panasonic",114 "booking.com", "expedia", "airbnb", "tripadvisor", "yelp",115 "doordash", "ubereats", "grubhub", "lyft", "uber",116 "twitch", "discord", "telegram", "whatsapp", "snapchat", "tiktok",117]118 119# ── Entity lists for rules normalizer ──────────────────────────────────────────120 121# Common IATA codes (2-3 letter airline codes)122IATA_CODES = {123 "AA", "BA", "DL", "UA", "LH", "AF", "EK", "QR", "SQ", "CX",124 "VS", "KL", "IB", "TK", "AC", "QF", "NH", "JL", "MH", "TG",125 "AI", "SA", "ET", "KE", "OZ", "CI", "BR", "LA", "AV", "AM",126 "WN", "B6", "AS", "F9", "NK", "G4", "VX", "HA",127}128 129# Common stock tickers → company name aliases130STOCK_ALIASES: dict[str, list[str]] = {131 "AAPL": ["apple", "aapl"],132 "TSLA": ["tesla", "tsla"],133 "MSFT": ["microsoft", "msft"],134 "GOOGL": ["google", "alphabet", "googl"],135 "AMZN": ["amazon", "amzn"],136 "META": ["meta", "facebook", "fb"],137 "NVDA": ["nvidia", "nvda"],138 "NFLX": ["netflix", "nflx"],139 "PYPL": ["paypal", "pypl"],140 "SNAP": ["snapchat", "snap"],141 "AMD": ["amd"],142 "INTC": ["intel", "intc"],143 "UBER": ["uber"],144 "LYFT": ["lyft"],145 "ABNB": ["airbnb", "abnb"],146 "COIN": ["coinbase", "coin"],147 "HOOD": ["robinhood", "hood"],148}149 150# Reverse map: alias → ticker151_ALIAS_TO_TICKER: dict[str, str] = {}152for ticker, aliases in STOCK_ALIASES.items():153 for alias in aliases:154 _ALIAS_TO_TICKER[alias.lower()] = ticker155 156# Product model patterns: brand → canonical prefix157PRODUCT_BRANDS = ["iphone", "samsung", "macbook", "ipad", "pixel", "surface"]158 159# ── Base normalizer ────────────────────────────────────────────────────────────160 161class Normalizer(ABC):162 name: str163 164 def warmup(self) -> None:165 """Called once before benchmarking to initialize any lazy state."""166 pass167 168 @abstractmethod169 def normalize(self, query: str) -> str:170 ...171 172 def normalize_batch(self, queries: list[str]) -> list[str]:173 return [self.normalize(q) for q in queries]174 175 176# ── 1. Identity (baseline) ────────────────────────────────────────────────────177 178class IdentityNormalizer(Normalizer):179 name = "Identity (baseline)"180 181 def normalize(self, query: str) -> str:182 return query183 184 185# ── 2. PySpellChecker ────────────────────────────────────────────────────────186 187class PySpellNormalizer(Normalizer):188 name = "PySpellChecker"189 190 def __init__(self):191 if not HAS_PYSPELL:192 raise RuntimeError("pyspellchecker not installed")193 self._sc = _SC()194 195 def normalize(self, query: str) -> str:196 words = query.lower().split()197 return " ".join(self._sc.correction(w) or w for w in words)198 199 200# ── 3. SymSpell ───────────────────────────────────────────────────────────────201 202_ORCAS_VOCAB = Path(__file__).parent / "orcas_vocab.txt"203 204 205class SymSpellNormalizer(Normalizer):206 name = "SymSpell"207 208 def __init__(self, max_edit_distance: int = 2):209 if not HAS_SYMSPELL:210 raise RuntimeError("symspellpy not installed")211 self._sym = _SS(max_dictionary_edit_distance=max_edit_distance)212 # Try importlib.resources first (works in newer Python/packaging setups),213 # fall back to pkg_resources for older environments.214 _dict_loaded = False215 # Try candidate dictionary filenames (name changed across symspellpy versions)216 _DICT_CANDIDATES = ["frequency_dictionary_en_82_765.txt", "en-80k.txt"]217 try:218 import importlib.resources as _ir219 for _fname in _DICT_CANDIDATES:220 try:221 _ref = _ir.files("symspellpy").joinpath(_fname)222 with _ir.as_file(_ref) as _dp:223 _dict_loaded = self._sym.load_dictionary(str(_dp), term_index=0, count_index=1)224 if _dict_loaded:225 break226 except Exception:227 pass228 except Exception:229 pass230 if not _dict_loaded:231 for _fname in _DICT_CANDIDATES:232 _dp = pkg_resources.resource_filename("symspellpy", _fname)233 _dict_loaded = self._sym.load_dictionary(_dp, term_index=0, count_index=1)234 if _dict_loaded:235 break236 if _ORCAS_VOCAB.exists():237 self._sym.load_dictionary(str(_ORCAS_VOCAB), term_index=0, count_index=1)238 self.name = "SymSpell+ORCAS"239 self._max_ed = max_edit_distance240 241 def normalize(self, query: str) -> str:242 # Use lookup_compound for multi-token correction243 suggestions = self._sym.lookup_compound(244 query.lower(), max_edit_distance=self._max_ed245 )246 if suggestions:247 return suggestions[0].term248 return query.lower()249 250 251# ── 4. Rules (entity + regex) ────────────────────────────────────────────────252 253class RulesNormalizer(Normalizer):254 name = "Rules (entity + regex)"255 256 # Flight: digits + IATA or IATA + digits → IATA + digits (no space)257 _FLIGHT_LOOSE = re.compile(258 r'\b(?:flight\s+)?(\d{2,4})\s*([A-Z]{2,3})\b' # 163 SQ259 r'|'260 r'\b(?:flight\s+)?([A-Z]{2,3})\s+(\d{2,4})\b', # SQ 163 (space)261 re.IGNORECASE262 )263 264 # Product spacing: brand directly followed by digits/variant ("iphone15")265 _PRODUCT_SPACING = re.compile(266 r'\b(iphone|macbook|ipad|pixel|galaxy|surface|airpods)'267 r'(\d+|pro|air|mini|max|ultra|plus)\b',268 re.IGNORECASE269 )270 271 # Stock: remove surrounding noise, keep just the ticker272 _STOCK_NOISE = re.compile(273 r'\b(stock|share|price|shares|equity|ticker|market|trading|invest(?:ment)?)\b',274 re.IGNORECASE275 )276 277 # Common compound words that users type without a space.278 # Applied per-token so works in multi-token queries too279 # e.g. "restarants nearme" → "restarants near me" (then GuardedPySpell fixes "restarants")280 _COMPOUND_SPLITS: dict[str, str] = {281 "nearme": "near me",282 "nearbyme": "near by me",283 "newyork": "new york",284 "losangeles": "los angeles",285 "sanfrancisco": "san francisco",286 "lasvegас": "las vegas",287 "bestbuy": "best buy",288 "homedepot": "home depot",289 "wholefoods": "whole foods",290 "starbucks": "starbucks", # already one word, no-op291 "doordash": "doordash",292 "ubereats": "uber eats",293 "grubhub": "grubhub",294 "openai": "openai",295 "chatgpt": "chatgpt",296 "youtube": "youtube",297 "facebook": "facebook",298 "instagram": "instagram",299 "whatsapp": "whatsapp",300 "linkedin": "linkedin",301 "tiktok": "tiktok",302 }303 304 def _normalize_flight(self, query: str) -> str:305 q_upper = query.upper()306 def _repl(m):307 if m.group(1): # digits IATA308 num, code = m.group(1), m.group(2).upper()309 else: # IATA digits310 code, num = m.group(3).upper(), m.group(4)311 if code in IATA_CODES:312 return f"{code}{num}"313 return m.group(0)314 result = self._FLIGHT_LOOSE.sub(_repl, query)315 return result316 317 def _normalize_stock(self, query: str) -> Optional[str]:318 ql = query.lower().strip()319 tokens = ql.split()320 # Check if any token is a known ticker or alias321 found_ticker = None322 for tok in tokens:323 # Direct ticker match (uppercase)324 if tok.upper() in STOCK_ALIASES:325 found_ticker = tok.upper()326 break327 # Alias match328 if tok in _ALIAS_TO_TICKER:329 found_ticker = _ALIAS_TO_TICKER[tok]330 if found_ticker:331 # Case 1: stock noise words present (e.g. "AAPL stock price")332 remaining = self._STOCK_NOISE.sub("", ql).strip()333 if remaining != ql.strip():334 return found_ticker335 # Case 2: explicit ticker token present alongside alias336 # (e.g. "apple aapl", "google GOOGL") — but NOT "google pixel 8"337 if found_ticker.lower() in tokens:338 return found_ticker339 return None340 341 def _normalize_product_spacing(self, query: str) -> str:342 return self._PRODUCT_SPACING.sub(lambda m: f"{m.group(1)} {m.group(2)}", query)343 344 def _normalize_compounds(self, query: str) -> str:345 """Split known compound tokens anywhere in the query.346 Works per-token so handles mixed queries like 'restarants nearme'."""347 tokens = query.lower().split()348 return " ".join(self._COMPOUND_SPLITS.get(tok, tok) for tok in tokens)349 350 def _normalize_word_order(self, query: str) -> str:351 """Reorder product queries so the brand/product-line token comes first.352 353 Handles patterns like:354 's24 samsung' → 'samsung s24'355 'pro 14 macbook' → 'macbook pro 14'356 'ultra s23 samsung'→ 'samsung ultra s23'357 """358 tokens = query.lower().split()359 if len(tokens) < 2:360 return query361 # Find a PRODUCT_BRANDS token that is not already at position 0362 for i, tok in enumerate(tokens):363 if i > 0 and tok in PRODUCT_BRANDS:364 # Move brand to front, preserve relative order of the rest365 return " ".join([tok] + tokens[:i] + tokens[i + 1:])366 return query367 368 def normalize(self, query: str) -> str:369 q = query.strip()370 371 # 1. Stock canonicalization372 stock = self._normalize_stock(q)373 if stock:374 return stock375 376 # 2. Flight ID normalization377 q = self._normalize_flight(q)378 379 # 3. Compound splitting (nearme → near me, newyork → new york)380 q = self._normalize_compounds(q)381 382 # 4. Product spacing383 q = self._normalize_product_spacing(q)384 385 # 5. Product word order386 q = self._normalize_word_order(q)387 388 # 6. Clean up extra whitespace389 q = re.sub(r'\s+', ' ', q).strip()390 391 return q392 393 394# ── 5. RapidFuzz (brand matching) ────────────────────────────────────────────395 396class RapidFuzzNormalizer(Normalizer):397 name = "RapidFuzz (brand match)"398 399 def __init__(self, score_cutoff: int = 82):400 if not HAS_RAPIDFUZZ:401 raise RuntimeError("rapidfuzz not installed")402 self._cutoff = score_cutoff403 404 def normalize(self, query: str) -> str:405 ql = query.lower().strip()406 407 # Only attempt brand correction on short queries (≤ 3 tokens)408 tokens = ql.split()409 if len(tokens) > 3:410 return query411 412 # Skip very short queries — too ambiguous to fuzzy-match safely413 # (e.g. 'appl', 'npm', 'gcc' should not be matched to brand names)414 if len(ql) <= 5:415 return query416 417 # Try matching each n-gram of the query against the brand list418 # First try the full query, then try progressively smaller windows419 result = _rf_process.extractOne(420 ql, BRANDS,421 scorer=_rf_fuzz.token_sort_ratio,422 score_cutoff=self._cutoff,423 )424 if result:425 best_match, score, _ = result426 return best_match427 428 return query429 430 431# ── 6. Combined ───────────────────────────────────────────────────────────────432 433class CombinedNormalizer(Normalizer):434 name = "Combined (Rules + SymSpell + RapidFuzz)"435 436 def __init__(self):437 self._rules = RulesNormalizer()438 self._symspell = SymSpellNormalizer() if HAS_SYMSPELL else None439 self._rfuzz = RapidFuzzNormalizer() if HAS_RAPIDFUZZ else None440 441 def normalize(self, query: str) -> str:442 q = query.strip()443 444 # Step 1: Apply entity/structural rules first (highest precision)445 q_rules = self._rules.normalize(q)446 if q_rules.lower() != q.lower():447 return q_rules # Rules made a change — trust it448 449 # Step 2: SymSpell for general typo correction450 if self._symspell:451 q_sym = self._symspell.normalize(q)452 if q_sym.lower() != q.lower():453 return q_sym454 455 # Step 3: RapidFuzz for brand name typos (catches what SymSpell misses456 # on compound brand names like "bestbuyt" → "best buy")457 if self._rfuzz:458 q_rf = self._rfuzz.normalize(q)459 if q_rf.lower() != q.lower():460 return q_rf461 462 return q463 464 465# ── 7. GuardedPySpell ────────────────────────────────────────────────────────466 467class GuardedPySpellNormalizer(Normalizer):468 """PySpellChecker with guards to prevent over-correction.469 470 PySpellChecker gets 88% on single_typo and 71% on multi_typo, but has471 40% over-correction on no-change queries (e.g. 'appl' → 'apple').472 473 Guards:474 - Skip tokens ≤ 4 chars (appl, npm, gcc, css, java, rust, echo, go)475 - Skip all-uppercase tokens (AAPL, NYC, SQ — abbreviations/tickers)476 - Skip tokens in the brand allowlist (airbnb, spotify, linkedin, …)477 478 Most legitimate short abbreviations are ≤ 4 chars or all-caps.479 Typos worth correcting are almost always ≥ 5 chars ('wheather', 'suhsi').480 """481 name = "PySpell (guarded)"482 483 # Known brand/product tokens that PySpell would corrupt.484 # Stored lowercase; comparison is done after lowercasing the token.485 _BRAND_ALLOWLIST: frozenset = frozenset({486 # Social / streaming487 "airbnb", "spotify", "linkedin", "tiktok", "whatsapp",488 "snapchat", "pinterest", "twitch", "reddit", "tumblr",489 "discord", "telegram", "signal",490 # Tech / SaaS491 "github", "gitlab", "dropbox", "notion", "figma",492 "asana", "trello", "jira", "confluence", "zendesk",493 "hubspot", "salesforce", "shopify", "stripe", "twilio",494 "vercel", "netlify", "supabase", "kubernetes", "terraform",495 "ansible", "grafana", "splunk", "datadog", "snowflake",496 "databricks", "pytorch", "tensorflow", "sklearn",497 # Devices / brands498 "iphone", "ipad", "macbook", "airpods", "homepod",499 "samsung", "pixel", "oneplus", "lenovo", "thinkpad",500 "playstation", "nintendo", "xbox",501 # Services502 "doordash", "grubhub", "instacart", "postmates",503 "lyft", "ubereats",504 # Media505 "netflix", "hulu", "disney", "peacock", "paramount",506 "youtube", "twitch",507 # Finance508 "venmo", "paypal", "cashapp", "robinhood", "coinbase",509 # Misc tech terms PySpell corrupts510 "nginx", "kafka", "numpy", "pandas",511 })512 513 def __init__(self):514 if not HAS_PYSPELL:515 raise RuntimeError("pyspellchecker not installed")516 self._sc = _SC()517 518 def _skip(self, token: str) -> bool:519 return len(token) <= 4 or token.isupper() or token in self._BRAND_ALLOWLIST520 521 def normalize(self, query: str) -> str:522 words = query.lower().split()523 return " ".join(524 w if self._skip(w) else (self._sc.correction(w) or w)525 for w in words526 )527 528 529# ── 8. CombinedV2 (Rules + GuardedPySpell + RapidFuzz) ───────────────────────530 531class CombinedV2Normalizer(Normalizer):532 """Improved pipeline: Rules → RapidFuzz (single-token) → SymSpell split → GuardedPySpell → RapidFuzz (multi-token).533 534 Rules handles structured entities (flight IDs, stock tickers, product535 spacing/order) with perfect precision. RapidFuzz runs first on single-token536 queries to catch brand typos (bestbuyt→best buy) before SymSpell can corrupt537 them (bestbuyt→best but). SymSpell compound splitting then handles concatenated538 words (nearme→near me). GuardedPySpell handles general typos while protecting539 short tokens. RapidFuzz runs again at the end for multi-token brand typos.540 """541 name = "CombinedV2 (Rules + GuardedPySpell + RapidFuzz)"542 543 def __init__(self):544 self._rules = RulesNormalizer()545 self._symspell = SymSpellNormalizer() if HAS_SYMSPELL else None546 self._pyspell = GuardedPySpellNormalizer() if HAS_PYSPELL else None547 self._rfuzz = RapidFuzzNormalizer() if HAS_RAPIDFUZZ else None548 549 def normalize(self, query: str) -> str:550 q = query.strip()551 552 # Step 1: Rules — flight IDs, stock tickers, product spacing/order553 q_rules = self._rules.normalize(q)554 if q_rules.lower() != q.lower():555 return q_rules556 557 # Step 2: RapidFuzz — brand name typos for single-token queries.558 # Must run before SymSpell compound splitting: SymSpell splits 'bestbuyt'559 # into 'best but' (wrong) whereas RapidFuzz correctly maps it to 'best buy'.560 if self._rfuzz and ' ' not in q:561 q_rf = self._rfuzz.normalize(q)562 if q_rf.lower() != q.lower():563 return q_rf564 565 # Step 3: SymSpell compound splitting for single-token queries only.566 # Only accept if SymSpell introduces a space (compound split).567 # Known compounds (nearme, newyork etc.) are handled by Rules above,568 # so this catches any remaining edge cases for single-token inputs.569 if self._symspell and ' ' not in q:570 q_sym = self._symspell.normalize(q)571 if ' ' in q_sym:572 return q_sym573 574 # Step 4: GuardedPySpell — general typos (skips short/uppercase tokens)575 if self._pyspell:576 q_spell = self._pyspell.normalize(q)577 if q_spell.lower() != q.lower():578 return q_spell579 580 # Step 5: RapidFuzz — brand name typos for multi-token queries581 # (e.g. 'gooogle maps' → 'google maps', 'spotifiy premium' → 'spotify premium')582 if self._rfuzz:583 q_rf = self._rfuzz.normalize(q)584 if q_rf.lower() != q.lower():585 return q_rf586 587 return q588 589 590# ── 9. ContextualSpellCheck (spaCy + BERT) ───────────────────────────────────591 592class ContextualSpellCheckNormalizer(Normalizer):593 """Uses BERT contextual embeddings to decide whether and how to correct594 each token. Unlike SymSpell, it sees the full query context before595 making a correction — so 'appl' in an ambiguous context stays as-is,596 while 'wheather nyc' correctly becomes 'weather nyc'.597 598 Requires:599 pip install contextualSpellCheck600 python -m spacy download en_core_web_sm601 """602 name = "ContextualSpellCheck (BERT)"603 604 def __init__(self):605 if not HAS_CONTEXTUAL:606 raise RuntimeError("contextualSpellCheck not available")607 self._nlp = _csc_nlp608 609 def normalize(self, query: str) -> str:610 doc = self._nlp(query)611 # doc._.outcome_spellCheck is the full corrected string612 result = doc._.outcome_spellCheck613 return result if result else query614 615 616# ── 8. T5 Spell Corrector (HuggingFace) ──────────────────────────────────────617 618class T5SpellCorrector(Normalizer):619 """Fine-tuned T5 model for spelling correction.620 Model: oliverguhr/spelling-correction-english-base621 622 This is a seq2seq model trained on noisy→clean sentence pairs.623 It handles multi-token typos, word order, and spacing better than624 dictionary-based approaches, but at significantly higher latency.625 626 Expected latency: ~100–500ms on CPU, ~20–80ms on GPU.627 628 Requires:629 pip install transformers torch (or transformers sentencepiece)630 """631 name = "T5 (oliverguhr/spelling-correction)"632 633 _MODEL_ID = "oliverguhr/spelling-correction-english-base"634 635 def __init__(self):636 if not HAS_TRANSFORMERS:637 raise RuntimeError("transformers not installed")638 self._pipe = None # lazy load in warmup()639 640 def warmup(self) -> None:641 print(f" Loading {self._MODEL_ID}...", end=" ", flush=True)642 self._pipe = _hf_pipeline(643 "text2text-generation",644 model=self._MODEL_ID,645 tokenizer=self._MODEL_ID,646 )647 # Prime the model with a dummy query648 self._pipe("warmup query", max_length=64)649 print("ready")650 651 def normalize(self, query: str) -> str:652 if self._pipe is None:653 self.warmup()654 result = self._pipe(query, max_length=128, num_beams=4)655 return result[0]["generated_text"].strip()656 657 658# ── 9. CombinedML (Rules → T5) ───────────────────────────────────────────────659 660class CombinedMLNormalizer(Normalizer):661 """Best-of-both-worlds pipeline:662 1. Rules handle structured entity normalization (flight IDs, stock tickers,663 product model reordering) with zero latency and perfect precision.664 2. T5 handles everything else — general typos, multi-token corrections,665 brand names — using full-query context.666 667 This avoids running T5 on queries that rules already handle perfectly,668 saving latency on the most common structured patterns.669 """670 name = "CombinedML (Rules → T5)"671 672 def __init__(self):673 self._rules = RulesNormalizer()674 self._t5 = T5SpellCorrector() if HAS_TRANSFORMERS else None675 676 def warmup(self) -> None:677 if self._t5:678 self._t5.warmup()679 680 def normalize(self, query: str) -> str:681 # Step 1: Rules first — highest precision for structured entities682 q_rules = self._rules.normalize(query)683 if q_rules.lower() != query.lower():684 return q_rules685 686 # Step 2: T5 for everything else687 if self._t5:688 return self._t5.normalize(query)689 690 return query691 692 693# ── Metrics ───────────────────────────────────────────────────────────────────694 695def char_error_rate(pred: str, gold: str) -> float:696 """CER = edit_distance / max(len(pred), len(gold))."""697 if not pred and not gold:698 return 0.0699 return edit_distance(pred.lower(), gold.lower()) / max(len(pred), len(gold))700 701 702def word_error_rate(pred: str, gold: str) -> float:703 """WER = token-level edit distance / number of gold tokens."""704 pred_toks = pred.lower().split()705 gold_toks = gold.lower().split()706 if not gold_toks:707 return 0.0708 m, n = len(pred_toks), len(gold_toks)709 dp = list(range(n + 1))710 for i in range(1, m + 1):711 prev = dp[:]712 dp[0] = i713 for j in range(1, n + 1):714 dp[j] = prev[j-1] if pred_toks[i-1] == gold_toks[j-1] \715 else 1 + min(prev[j], dp[j-1], prev[j-1])716 return dp[n] / n717 718 719def run_benchmark(normalizer: Normalizer, df: pd.DataFrame, n_timing_reps: int = 5) -> dict:720 """Run a normalizer on the dataset and return metrics."""721 queries = df["noisy"].tolist()722 723 # ── Timing ───────────────────────────────────────────────────────────────724 latencies_ms = []725 for q in queries:726 t0 = time.perf_counter()727 for _ in range(n_timing_reps):728 normalizer.normalize(q)729 t1 = time.perf_counter()730 latencies_ms.append((t1 - t0) / n_timing_reps * 1000)731 732 # ── Predictions ──────────────────────────────────────────────────────────733 preds = [normalizer.normalize(q) for q in queries]734 df = df.copy()735 df["pred"] = preds736 737 def em(row): return row["pred"].lower().strip() == row["canonical"].lower().strip()738 def cer(row): return char_error_rate(row["pred"], row["canonical"])739 def wer(row): return word_error_rate(row["pred"], row["canonical"])740 741 df["em"] = df.apply(em, axis=1)742 df["cer"] = df.apply(cer, axis=1)743 df["wer"] = df.apply(wer, axis=1)744 745 # No-change precision and over-correction rate746 nc = df[~df["should_change"]]747 no_change_precision = (nc["pred"].str.lower().str.strip() == nc["noisy"].str.lower().str.strip()).mean() if len(nc) else float("nan")748 over_correction = 1.0 - no_change_precision if not np.isnan(no_change_precision) else float("nan")749 750 # ── Per-category exact match ──────────────────────────────────────────────751 cat_em = df.groupby("category")["em"].mean().to_dict()752 753 return {754 "name": normalizer.name,755 "exact_match": df["em"].mean(),756 "cer_mean": df["cer"].mean(),757 "wer_mean": df["wer"].mean(),758 "no_change_precision": no_change_precision,759 "over_correction": over_correction,760 "latency_mean_ms": np.mean(latencies_ms),761 "latency_p50_ms": np.percentile(latencies_ms, 50),762 "latency_p95_ms": np.percentile(latencies_ms, 95),763 "latency_p99_ms": np.percentile(latencies_ms, 99),764 "per_category": cat_em,765 "_df": df, # store for detailed output766 "_latencies": latencies_ms,767 }768 769 770# ── Main ──────────────────────────────────────────────────────────────────────771 772def main():773 parser = argparse.ArgumentParser()774 parser.add_argument("--dataset", default=str(Path(__file__).parent / "dataset.csv"))775 parser.add_argument("--reps", type=int, default=5, help="Timing repetitions per query")776 args = parser.parse_args()777 778 df = pd.read_csv(args.dataset)779 print(f"Loaded {len(df)} rows from {args.dataset}")780 print(f"Categories: {df['category'].value_counts().to_dict()}\n")781 782 # ── Build normalizer list ─────────────────────────────────────────────────783 normalizers: list[Normalizer] = [IdentityNormalizer(), RulesNormalizer()]784 if HAS_PYSPELL:785 normalizers.append(PySpellNormalizer())786 if HAS_SYMSPELL:787 normalizers.append(SymSpellNormalizer())788 if HAS_RAPIDFUZZ:789 normalizers.append(RapidFuzzNormalizer())790 if HAS_SYMSPELL and HAS_RAPIDFUZZ:791 normalizers.append(CombinedNormalizer())792 if HAS_PYSPELL:793 normalizers.append(GuardedPySpellNormalizer())794 if HAS_PYSPELL and HAS_RAPIDFUZZ:795 normalizers.append(CombinedV2Normalizer())796 # ML normalizers (disabled — too slow and underperform rules-based)797 # if HAS_CONTEXTUAL:798 # normalizers.append(ContextualSpellCheckNormalizer())799 # if HAS_TRANSFORMERS:800 # normalizers.append(T5SpellCorrector())801 # normalizers.append(CombinedMLNormalizer())802 803 # Warmup804 for norm in normalizers:805 norm.warmup()806 807 # ── Run benchmarks ────────────────────────────────────────────────────────808 results = []809 for norm in normalizers:810 print(f"Benchmarking: {norm.name}...", end=" ", flush=True)811 r = run_benchmark(norm, df, n_timing_reps=args.reps)812 results.append(r)813 print(f"EM={r['exact_match']:.1%} CER={r['cer_mean']:.3f} lat_p50={r['latency_p50_ms']:.2f}ms")814 815 # ── Summary table ─────────────────────────────────────────────────────────816 print("\n" + "="*90)817 print("SUMMARY — Overall Metrics")818 print("="*90)819 820 summary_rows = []821 for r in results:822 summary_rows.append({823 "Normalizer": r["name"],824 "Exact Match": f"{r['exact_match']:.1%}",825 "CER": f"{r['cer_mean']:.3f}",826 "WER": f"{r['wer_mean']:.3f}",827 "No-change Prec.": f"{r['no_change_precision']:.1%}" if not np.isnan(r['no_change_precision']) else "N/A",828 "Over-correction": f"{r['over_correction']:.1%}" if not np.isnan(r['over_correction']) else "N/A",829 "Lat mean (ms)": f"{r['latency_mean_ms']:.2f}",830 "Lat p50 (ms)": f"{r['latency_p50_ms']:.2f}",831 "Lat p95 (ms)": f"{r['latency_p95_ms']:.2f}",832 "Lat p99 (ms)": f"{r['latency_p99_ms']:.2f}",833 })834 835 try:836 from tabulate import tabulate837 print(tabulate(summary_rows, headers="keys", tablefmt="rounded_outline"))838 except ImportError:839 pd.DataFrame(summary_rows).to_string(index=False)840 print(pd.DataFrame(summary_rows).to_string(index=False))841 842 # ── Per-category table ────────────────────────────────────────────────────843 categories = sorted(df["category"].unique())844 print("\n" + "="*90)845 print("PER-CATEGORY Exact Match")846 print("="*90)847 848 cat_rows = []849 for r in results:850 row = {"Normalizer": r["name"][:30]}851 for cat in categories:852 row[cat] = f"{r['per_category'].get(cat, float('nan')):.0%}"853 cat_rows.append(row)854 855 try:856 from tabulate import tabulate857 print(tabulate(cat_rows, headers="keys", tablefmt="rounded_outline"))858 except ImportError:859 print(pd.DataFrame(cat_rows).to_string(index=False))860 861 # ── Sample predictions ────────────────────────────────────────────────────862 print("\n" + "="*90)863 print("SAMPLE PREDICTIONS — Combined vs Identity (first 5 per category)")864 print("="*90)865 866 combined_r = next((r for r in results if "CombinedV2" in r["name"]),867 next((r for r in results if "Combined" in r["name"]), results[-1]))868 identity_r = results[0]869 870 for cat in categories:871 sub = combined_r["_df"][combined_r["_df"]["category"] == cat].head(5)872 id_sub = identity_r["_df"][identity_r["_df"]["category"] == cat].head(5)873 print(f"\n {cat.upper()}")874 print(f" {'Noisy':<30} {'Canonical':<25} {'Combined pred':<25} {'EM':>4}")875 print(f" {'-'*30} {'-'*25} {'-'*25} {'-'*4}")876 for (_, row), (_, id_row) in zip(sub.iterrows(), id_sub.iterrows()):877 em_mark = "✓" if row["em"] else "✗"878 print(f" {row['noisy']:<30} {row['canonical']:<25} {row['pred']:<25} {em_mark:>4}")879 880 # ── Save full results ─────────────────────────────────────────────────────881 out_path = Path(args.dataset).parent / "results.csv"882 combined_r["_df"].to_csv(out_path, index=False)883 print(f"\nFull predictions saved to {out_path}")884 885 886if __name__ == "__main__":887 main()888