Files
test-search-engine/app/search_index.py
T

259 lines
9.7 KiB
Python

from __future__ import annotations
import math
from collections import Counter, defaultdict
from dataclasses import dataclass
from typing import TYPE_CHECKING
from app.normalization import meaningful_tokens, normalize_text, tokenize, word_tokens
from app.typo import (
character_trigrams,
damerau_levenshtein_distance,
trigram_dice,
word_similarity,
)
if TYPE_CHECKING:
from app.catalog import CatalogItem
@dataclass(frozen=True, slots=True)
class LexicalScore:
item_index: int
score: float
bm25: float
trigram: float
token_alignment: float
exact_coverage: float
@dataclass(frozen=True, slots=True)
class ProductTypeInference:
product_type: str
score: float
query_word: str
catalog_word: str
class LexicalIndex:
"""Small in-memory lexical index: BM25 + char trigrams + typo alignment."""
def __init__(self, items: tuple[CatalogItem, ...]) -> None:
self._items = items
document_tokens: list[tuple[str, ...]] = []
document_words: list[tuple[str, ...]] = []
term_frequencies: list[Counter[str]] = []
document_frequency: Counter[str] = Counter()
trigram_postings: dict[str, set[int]] = defaultdict(set)
for index, item in enumerate(items):
text = f"{item.normalized_name} {item.product_type}"
tokens = meaningful_tokens(tokenize(text, already_normalized=True))
words = word_tokens(tokens)
frequencies = Counter(tokens)
document_tokens.append(tokens)
document_words.append(words)
term_frequencies.append(frequencies)
document_frequency.update(frequencies.keys())
for trigram in character_trigrams(" ".join(tokens)):
trigram_postings[trigram].add(index)
self._document_tokens = tuple(document_tokens)
self._document_words = tuple(document_words)
self._term_frequencies = tuple(term_frequencies)
self._document_lengths = tuple(len(tokens) for tokens in document_tokens)
self._average_document_length = max(
1.0, sum(self._document_lengths) / max(1, len(self._document_lengths))
)
document_count = len(items)
self._idf = {
term: math.log(1.0 + (document_count - frequency + 0.5) / (frequency + 0.5))
for term, frequency in document_frequency.items()
}
self._trigram_postings = {
trigram: frozenset(indexes) for trigram, indexes in trigram_postings.items()
}
def trigram_candidates(self, query_text: str, limit: int = 80) -> tuple[int, ...]:
normalized = normalize_text(query_text)
query_tokens = meaningful_tokens(tokenize(normalized, already_normalized=True))
counts: Counter[int] = Counter()
for trigram in character_trigrams(" ".join(query_tokens)):
counts.update(self._trigram_postings.get(trigram, ()))
return tuple(index for index, _ in counts.most_common(limit))
def infer_product_type(self, query_text: str) -> ProductTypeInference | None:
"""Infer a missing product type from stable lexical evidence.
Trigrams provide the shortlist, but are deliberately not sufficient on
their own: sizes and random strings can share a few character triples
with unrelated catalog rows. The leading rows must agree on one type
and share a catalog word that is an exact match, a morphological prefix,
or a tightly bounded typo of a query word.
"""
candidate_indices = self.trigram_candidates(query_text, limit=60)
if len(candidate_indices) < 3:
return None
leaders = self.rank(query_text, "", candidate_indices)[:3]
if len(leaders) < 3:
return None
product_types = {
self._items[leader.item_index].product_type for leader in leaders
}
if len(product_types) != 1:
return None
common_words = set(self._document_words[leaders[0].item_index])
for leader in leaders[1:]:
common_words.intersection_update(self._document_words[leader.item_index])
if not common_words:
return None
normalized = normalize_text(query_text)
query_words = word_tokens(
meaningful_tokens(tokenize(normalized, already_normalized=True))
)
best: tuple[float, str, str] | None = None
for query_word in query_words:
if len(query_word) < 4:
continue
for catalog_word in common_words:
score = self._type_inference_word_score(query_word, catalog_word)
if score <= 0:
continue
evidence = (score, query_word, catalog_word)
if best is None or evidence > best:
best = evidence
if best is None:
return None
score, query_word, catalog_word = best
return ProductTypeInference(
product_type=next(iter(product_types)),
score=score,
query_word=query_word,
catalog_word=catalog_word,
)
@staticmethod
def _type_inference_word_score(query_word: str, catalog_word: str) -> float:
if query_word == catalog_word:
return 1.0
shorter_length = min(len(query_word), len(catalog_word))
longer_length = max(len(query_word), len(catalog_word))
if shorter_length >= 4 and (
query_word.startswith(catalog_word)
or catalog_word.startswith(query_word)
):
length_ratio = shorter_length / longer_length
if length_ratio >= 0.55:
return 0.72 + length_ratio * 0.18
distance = damerau_levenshtein_distance(query_word, catalog_word)
if shorter_length >= 5 and distance <= 1:
return 0.90
if longer_length >= 8 and distance <= 2:
edit_similarity = 1.0 - distance / longer_length
if edit_similarity >= 0.75:
return 0.82
return 0.0
def rank(
self,
query_text: str,
product_type: str,
candidate_indices: tuple[int, ...],
) -> list[LexicalScore]:
normalized = normalize_text(query_text)
query_tokens = list(meaningful_tokens(tokenize(normalized, already_normalized=True)))
# Alias normalization adds the canonical type as a lexical signal but
# leaves the original wording intact for typo evidence.
query_tokens.extend(meaningful_tokens(tokenize(product_type)))
query_token_tuple = tuple(query_tokens)
query_words = word_tokens(query_token_tuple)
query_counter = Counter(query_token_tuple)
query_trigram_text = " ".join(query_token_tuple)
raw_bm25: dict[int, float] = {}
for item_index in candidate_indices:
raw_bm25[item_index] = self._bm25(item_index, query_counter)
maximum_bm25 = max(raw_bm25.values(), default=0.0)
scores: list[LexicalScore] = []
for item_index in candidate_indices:
item_tokens = self._document_tokens[item_index]
item_words = self._document_words[item_index]
bm25 = raw_bm25[item_index] / maximum_bm25 if maximum_bm25 > 0 else 0.0
trigram = trigram_dice(query_trigram_text, " ".join(item_tokens))
alignment = self._token_alignment(query_words, item_words)
exact = self._exact_coverage(query_token_tuple, item_tokens)
combined = bm25 * 0.42 + trigram * 0.23 + alignment * 0.25 + exact * 0.10
scores.append(
LexicalScore(
item_index=item_index,
score=max(0.0, min(1.0, combined)),
bm25=bm25,
trigram=trigram,
token_alignment=alignment,
exact_coverage=exact,
)
)
scores.sort(key=lambda item: (-item.score, self._items[item.item_index].sku))
return scores
def _bm25(self, item_index: int, query: Counter[str]) -> float:
k1 = 1.5
b = 0.75
frequencies = self._term_frequencies[item_index]
document_length = self._document_lengths[item_index]
score = 0.0
for term, query_frequency in query.items():
term_frequency = frequencies.get(term, 0)
if term_frequency == 0:
continue
denominator = term_frequency + k1 * (
1.0 - b + b * document_length / self._average_document_length
)
score += (
self._idf.get(term, 0.0)
* term_frequency
* (k1 + 1.0)
/ denominator
* min(query_frequency, 2)
)
return score
@staticmethod
def _token_alignment(
query_words: tuple[str, ...], document_words: tuple[str, ...]
) -> float:
if not query_words:
return 0.0
values: list[float] = []
for query_word in query_words:
if query_word in document_words:
values.append(1.0)
continue
if len(query_word) < 4:
values.append(0.0)
continue
best = max(
(word_similarity(query_word, document_word) for document_word in document_words),
default=0.0,
)
values.append(best if best >= 0.50 else 0.0)
return sum(values) / len(values)
@staticmethod
def _exact_coverage(query_tokens: tuple[str, ...], document_tokens: tuple[str, ...]) -> float:
if not query_tokens:
return 0.0
document_set = set(document_tokens)
return sum(token in document_set for token in query_tokens) / len(query_tokens)