From b30281d159625d54aa917ab5d1282327e05cd081 Mon Sep 17 00:00:00 2001 From: Alan Scherger Date: Mon, 19 Oct 2026 09:07:25 -0500 Subject: [PATCH] feat: Lightly stem BM25 tokens Natural-language queries often use a different inflection than the code they target ("handlers" vs EventHandler, "parsing" vs parse). Strip plurals, -ing and -ed from every BM25 token, at index and query time alike. Short tokens and tokens with digits or underscores are left untouched. Bump CACHE_FORMAT_VERSION so existing indexes are rebuilt with the new tokens. --- src/semble/index/types.py | 2 +- src/semble/tokens.py | 31 +++++++++++++++++++++++++++++-- tests/test_tokens.py | 32 ++++++++++++++++++++++++++++++++ 3 files changed, 62 insertions(+), 3 deletions(-) create mode 100644 tests/test_tokens.py diff --git a/src/semble/index/types.py b/src/semble/index/types.py index a08b7b6d3..6bf173e41 100644 --- a/src/semble/index/types.py +++ b/src/semble/index/types.py @@ -6,7 +6,7 @@ from semble.index.bm25 import BM25 from semble.types import Chunk, EmbeddingMatrix -CACHE_FORMAT_VERSION = 2 # Bump when the persisted index schema changes. +CACHE_FORMAT_VERSION = 3 # Bump when the persisted index schema or BM25 tokenization changes. def make_chunk_id(indexed_path: str, slot: int) -> str: diff --git a/src/semble/tokens.py b/src/semble/tokens.py index ab744e91a..3e2ed8c37 100644 --- a/src/semble/tokens.py +++ b/src/semble/tokens.py @@ -8,6 +8,9 @@ # "XMLParser" -> ["XML", "Parser"] _CAMEL_RE = re.compile(r"[A-Z]+(?=[A-Z][a-z])|[A-Z]?[a-z]+|[A-Z]+|[0-9]+") +# Plural "-es" after a sibilant: "classes" -> "class", "matches" -> "match". +_ES_AFTER_SIBILANT_RE = re.compile(r"(?:s|x|z|ch|sh)es$") + def split_identifier(token: str) -> list[str]: """Split a single identifier into sub-tokens via camelCase/snake_case. @@ -32,15 +35,39 @@ def split_identifier(token: str) -> list[str]: return [lower] +def stem(token: str) -> str: + """Strip common English inflections so "handlers"/"handler" and "parsed"/"parse" share a BM25 term. + + Deliberately light (plurals, -ing, -ed) and applied to queries and documents alike, + so over-stemming only merges terms on both sides. Tokens of 3 characters or fewer, + and tokens with digits or underscores, are left as they are. + E.g. "handlers" -> "handler", "classes" -> "class", "dependencies" -> "dependency" + """ + if len(token) <= 3 or not token.isalpha(): + return token + if token.endswith("ies") and len(token) > 4: + return token[:-3] + "y" + if _ES_AFTER_SIBILANT_RE.search(token): + return token[:-2] + if token.endswith("s") and not token.endswith(("ss", "us", "is")): + return token[:-1] + if token.endswith("ing") and len(token) > 5: + return token[:-3] + if token.endswith("ed") and len(token) > 4: + return token[:-2] + return token + + def tokenize(text: str) -> list[str]: """Split text into lowercase identifier-like tokens for BM25 indexing. Compound identifiers (camelCase, PascalCase, snake_case) are expanded into sub-tokens so that partial matches work. The original compound - token is preserved for exact-match boosting. + token is preserved for exact-match boosting. Every token is passed + through :func:`stem`, so inflected forms share a term. """ raw_tokens = _TOKEN_RE.findall(text) result: list[str] = [] for tok in raw_tokens: - result.extend(split_identifier(tok)) + result.extend(stem(part) for part in split_identifier(tok)) return result diff --git a/tests/test_tokens.py b/tests/test_tokens.py new file mode 100644 index 000000000..4787eb19c --- /dev/null +++ b/tests/test_tokens.py @@ -0,0 +1,32 @@ +import pytest + +from semble.tokens import stem, tokenize + + +@pytest.mark.parametrize( + ("token", "expected"), + [ + ("handlers", "handler"), + ("classes", "class"), + ("matches", "match"), + ("dependencies", "dependency"), + ("parsed", "pars"), + ("parsing", "pars"), + ("status", "status"), # -us is not a plural + ("process", "process"), # -ss is not a plural + ("analysis", "analysis"), # -is is not a plural + ("ids", "ids"), # 3 characters or fewer: unchanged + ("http2", "http2"), # digits: unchanged + ("my_func", "my_func"), # underscores: unchanged + ], +) +def test_stem(token: str, expected: str) -> None: + """The stemmer strips plurals, -ing and -ed, and leaves short or non-alphabetic tokens alone.""" + assert stem(token) == expected + + +def test_tokenize_stems_query_and_document_alike() -> None: + """An inflected query word and its base form in code produce the same BM25 term.""" + assert "handler" in tokenize("registered handlers") + assert "handler" in tokenize("class EventHandler:") + assert tokenize("RequestHandlers") == ["requesthandler", "request", "handler"]