kkthakur's picture
Deploy Local Hybrid Engine
b336134
Raw
History Blame Contribute Delete
4.06 kB
"""
Spelling correction module using SymSpell (Symmetric Delete Spelling Correction).
Optimized for O(1) delete lookups to match column names and operation keywords.
"""
from __future__ import annotations
import re
def _get_deletes(word: str, max_edit_distance: int = 2) -> set[str]:
"""Generate deletes for a word up to a maximum edit distance."""
deletes = set()
queue = {word}
for _ in range(max_edit_distance):
next_queue = set()
for w in queue:
if len(w) > 1:
for i in range(len(w)):
del_w = w[:i] + w[i+1:]
deletes.add(del_w)
next_queue.add(del_w)
queue = next_queue
return deletes
def _lev_dist(s1: str, s2: str) -> int:
"""Calculate the Levenshtein distance between two strings."""
if len(s1) < len(s2):
return _lev_dist(s2, s1)
if len(s2) == 0:
return len(s1)
previous_row = list(range(len(s2) + 1))
for i, c1 in enumerate(s1):
current_row = [i + 1]
for j, c2 in enumerate(s2):
insertions = previous_row[j + 1] + 1
deletions = current_row[j] + 1
substitutions = previous_row[j] + (c1 != c2)
current_row.append(min(insertions, deletions, substitutions))
previous_row = current_row
return previous_row[-1]
class SymSpell:
"""A lightweight symmetric delete spelling corrector."""
def __init__(self, max_edit_distance: int = 2):
self.max_edit_distance = max_edit_distance
# Maps delete_item -> set of original_words
self.deletes: dict[str, set[str]] = {}
self.words: set[str] = set()
def add_word(self, word: str) -> None:
"""Index a word and its deletes for O(1) spelling lookup."""
word = word.lower().strip()
if not word or word in self.words:
return
self.words.add(word)
# Index word itself
if word not in self.deletes:
self.deletes[word] = set()
self.deletes[word].add(word)
# Index deletions
for delete in _get_deletes(word, self.max_edit_distance):
if delete not in self.deletes:
self.deletes[delete] = set()
self.deletes[delete].add(word)
def lookup(self, word: str) -> list[str]:
"""Find candidate words matching the spelling of input word."""
word = word.lower().strip()
if not word:
return []
if word in self.words:
return [word]
candidates: set[str] = set()
# 1. Direct delete match
if word in self.deletes:
candidates.update(self.deletes[word])
# 2. Deletes of word match
for delete in _get_deletes(word, self.max_edit_distance):
if delete in self.deletes:
candidates.update(self.deletes[delete])
if delete in self.words:
candidates.add(delete)
# Score and rank candidates by Levenshtein distance
scored = []
for cand in candidates:
dist = _lev_dist(word, cand)
if dist <= self.max_edit_distance:
scored.append((cand, dist))
# Sort by distance first, then length (longer words first for ties)
scored.sort(key=lambda x: (x[1], -len(x[0])))
return [c for c, _ in scored]
def correct_query(self, query: str) -> str:
"""Correct typos in words within the query string."""
# Find all alphabet-only words
words = re.findall(r'[a-zA-Z]+', query)
corrected = query
for w in words:
if len(w) > 2: # Only correct words longer than 2 characters
suggestions = self.lookup(w)
if suggestions:
# Match exact word boundary to prevent partial replacements
corrected = re.sub(rf'\b{re.escape(w)}\b', suggestions[0], corrected, flags=re.IGNORECASE)
return corrected