feat: Lightly stem BM25 tokens

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:

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

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"]