LLM example retrieval

machine/translation/llm/__init__.py

@@ -0,0 +1,17 @@
+from .coverage_example_retriever import CoverageExampleRetriever
+from .example import Example
+from .example_retriever_factory import ExampleRetrieverFactory
+from .example_retriever_method import ExampleRetrievalMethod
+from .lazy_example_retriever import LazyExampleRetriever
+from .tfidf_example_retriever import TfidfExampleRetriever
+from .top_k_example_retriever import TopKExampleRetriever
+
+__all__ = [
+    "CoverageExampleRetriever",
+    "Example",
+    "ExampleRetrieverFactory",
+    "ExampleRetrievalMethod",
+    "LazyExampleRetriever",
+    "TfidfExampleRetriever",
+    "TopKExampleRetriever",
+]

machine/translation/llm/coverage_example_retriever.py

@@ -0,0 +1,95 @@
+from typing import Iterable, List, Optional, cast
+
+import numpy as np
+from scipy.sparse import csr_matrix
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import linear_kernel
+
+from .example import Example
+from .lazy_example_retriever import LazyExampleRetriever
+
+
+class CoverageExampleRetriever(LazyExampleRetriever):
+    # Benchmarked as a set; the measured gain says nothing about other values.
+    _ANCHORS = 20
+    _CANDIDATE_MULTIPLE = 4
+    _MIN_CANDIDATES = 100
+    _RELEVANCE_WEIGHT = 0.7
+
+    def __init__(self, examples: Iterable[Example]) -> None:
+        super().__init__(examples)
+        self._vectorizer: Optional[TfidfVectorizer] = None
+        self._matrix: Optional[csr_matrix] = None
+
+    def _fit(self) -> None:
+        # TfidfVectorizer rejects a corpus with nothing to put in its vocabulary.
+        sources = [ex.source for ex in self.examples]
+        if not any(source.strip() for source in sources):
+            self._vectorizer = None
+            self._matrix = None
+            return
+        self._vectorizer = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 5))
+        self._matrix = cast(csr_matrix, self._vectorizer.fit_transform(sources))
+
+    def _rank(self, query: str, k: int, exclude_index: Optional[int] = None) -> List[Example]:
+        if self._vectorizer is None or self._matrix is None:
+            return []
+        vector = cast(csr_matrix, self._vectorizer.transform([query]))
+        similarities = linear_kernel(vector, self._matrix).ravel()
+        ranked = similarities.argsort()[::-1]
+        if exclude_index is not None:
+            ranked = ranked[ranked != exclude_index]
+        k = min(k, len(ranked))
+        if k == 0:
+            return []
+        if vector.nnz == 0:
+            return [self.examples[i] for i in ranked[:k].tolist()]
+        candidates = self._distinct_candidates(
+            ranked, k, None if exclude_index is None else self.examples[exclude_index]
+        )
+        if len(candidates) == 0:
+            return []
+        chosen = self._choose_for_coverage(vector, self._matrix[candidates], similarities[candidates], k)
+        # Ties are broken by position so that the same corpus always yields the same prompt.
+        return [
+            self.examples[example_index]
+            for example_index in sorted(
+                (int(candidates[i]) for i in chosen), key=lambda i: (similarities[i], i), reverse=True
+            )
+        ]
+
+    def _distinct_candidates(self, ranked: np.ndarray, k: int, excluded: Optional[Example]) -> np.ndarray:
+        limit = max(self._CANDIDATE_MULTIPLE * k, self._MIN_CANDIDATES)
+        candidates: List[int] = []
+        # A copy of the excluded example would show the answer to the row it is excluded for.
+        seen = set() if excluded is None else {excluded}
+        for index in ranked:
+            example = self.examples[index]
+            if example not in seen:
+                candidates.append(int(index))
+                seen.add(example)
+            if len(candidates) >= limit:
+                break
+        return np.asarray(candidates)
+
+    def _choose_for_coverage(self, vector, candidate_matrix, similarities: np.ndarray, k: int) -> List[int]:
+        features = candidate_matrix[:, vector.indices].toarray() > 0
+        weights = vector.data / vector.data.sum()
+        relevance = similarities / max(float(similarities.max()), 1e-12)
+        # The closest few are kept outright, so that strong whole-sentence analogues survive the balancing.
+        anchor_count = min(self._ANCHORS, max(1, k // 2), len(features))
+        chosen = list(range(anchor_count))
+        available = np.ones(len(features), dtype=bool)
+        available[chosen] = False
+        counts = features[chosen].sum(axis=0).astype(float)
+        while len(chosen) < min(k, len(features)):
+            # Each n-gram counts for less the more of the chosen examples already demonstrate it.
+            gain = features @ (weights / (1.0 + counts))
+            gain /= max(float(gain[available].max()), 1e-12)
+            utility = self._RELEVANCE_WEIGHT * relevance + (1 - self._RELEVANCE_WEIGHT) * gain
+            utility[~available] = -np.inf
+            best = int(np.argmax(utility))
+            chosen.append(best)
+            available[best] = False
+            counts += features[best]
+        return chosen

machine/translation/llm/example.py

@@ -0,0 +1,7 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Example:
+    source: str
+    target: str

machine/translation/llm/example_retriever_factory.py

@@ -0,0 +1,19 @@
+from typing import Iterable
+
+from .coverage_example_retriever import CoverageExampleRetriever
+from .example import Example
+from .example_retriever_method import ExampleRetrievalMethod
+from .lazy_example_retriever import LazyExampleRetriever
+from .tfidf_example_retriever import TfidfExampleRetriever
+
+
+class ExampleRetrieverFactory:
+    @staticmethod
+    def create_example_retriever(method: ExampleRetrievalMethod, examples: Iterable[Example]) -> LazyExampleRetriever:
+        match method:
+            case ExampleRetrievalMethod.TFIDF:
+                return TfidfExampleRetriever(examples)
+            case ExampleRetrievalMethod.COVERAGE:
+                return CoverageExampleRetriever(examples)
+            case _:
+                raise ValueError(f"Unable to create example retriever with method {method}")

machine/translation/llm/example_retriever_method.py

@@ -0,0 +1,6 @@
+from enum import Enum, auto
+
+
+class ExampleRetrievalMethod(Enum):
+    TFIDF = auto()
+    COVERAGE = auto()

machine/translation/llm/lazy_example_retriever.py

@@ -0,0 +1,43 @@
+import threading
+from abc import ABC, abstractmethod
+from typing import Iterable, List, Optional, Union
+
+from .example import Example
+
+
+class LazyExampleRetriever(ABC):
+
+    def __init__(self, examples: Iterable[Example]) -> None:
+        self._examples = examples
+        self._fitted = False
+        self._lock = threading.RLock()
+
+    @property
+    def examples(self) -> List[Example]:
+        with self._lock:
+            if not isinstance(self._examples, List):
+                self._examples = list(self._examples)
+            return self._examples
+
+    @abstractmethod
+    def _fit(self) -> None: ...
+
+    def rank(self, query: str, k: int, exclude: Optional[Union[int, Example]] = None) -> List[Example]:
+        if k <= 0:
+            return []
+        num_examples = len(self.examples)
+        exclude_index = (
+            exclude
+            if isinstance(exclude, int) or exclude is None
+            else (self.examples.index(exclude) if exclude in self.examples else None)
+        )
+        if k >= num_examples:
+            return [example for index, example in enumerate(self.examples) if index != exclude_index]
+        with self._lock:
+            if not self._fitted:
+                self._fit()
+                self._fitted = True
+        return self._rank(query, k, exclude_index)
+
+    @abstractmethod
+    def _rank(self, query: str, k: int, exclude_index: Optional[int] = None) -> List[Example]: ...

machine/translation/llm/tfidf_example_retriever.py

@@ -0,0 +1,44 @@
+from typing import Iterable, List, Optional, cast
+
+import numpy as np
+from scipy.sparse import csr_matrix
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+from ...tokenization.latin_word_tokenizer import LatinWordTokenizer
+from .example import Example
+from .top_k_example_retriever import TopKExampleRetriever
+
+
+class TfidfExampleRetriever(TopKExampleRetriever):
+
+    def __init__(self, examples: Iterable[Example]) -> None:
+        super().__init__(examples)
+        self._latin_tokenizer = LatinWordTokenizer()
+        self._vectorizer: Optional[TfidfVectorizer] = None
+        self._matrix: Optional[csr_matrix] = None
+
+    def _fit(self) -> None:
+        # TfidfVectorizer rejects a corpus with nothing to put in its vocabulary.
+        sources = [ex.source for ex in self.examples]
+        if not any(self._tokenize(source) for source in sources):
+            self._vectorizer = None
+            self._matrix = None
+            return
+        self._vectorizer = TfidfVectorizer(
+            lowercase=False, tokenizer=self._tokenize, token_pattern=None  # pyright: ignore[reportArgumentType]
+        )
+        self._matrix = cast(csr_matrix, self._vectorizer.fit_transform(sources))
+
+    def _score_examples(self, query: str) -> np.ndarray:
+        if self._matrix is None or self._vectorizer is None:
+            return np.array([])
+        index_of_query = next((i for i, ex in enumerate(self.examples) if ex.source == query), -1)
+        if index_of_query != -1:
+            vector: csr_matrix = self._matrix[index_of_query]
+        else:
+            vector = cast(csr_matrix, self._vectorizer.transform([query]))
+        return (self._matrix @ vector.T).toarray().ravel()
+
+    def _tokenize(self, text: str) -> List[str]:
+        # Strip all punctuation-only tokens
+        return [token for token in self._latin_tokenizer.tokenize(text.lower()) if any(c.isalnum() for c in token)]

machine/translation/llm/top_k_example_retriever.py

@@ -0,0 +1,27 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+import numpy as np
+
+from machine.translation.llm.example import Example
+
+from .lazy_example_retriever import LazyExampleRetriever
+
+
+class TopKExampleRetriever(LazyExampleRetriever, ABC):
+    @abstractmethod
+    def _score_examples(self, query: str) -> np.ndarray: ...
+
+    def _rank(self, query: str, k: int, exclude_index: Optional[int] = None) -> List[Example]:
+        scores = self._score_examples(query)
+        n = scores.shape[0]
+        available = n - 1 if exclude_index is not None else n
+        k = min(k, available)
+        if k <= 0:
+            return []
+        if exclude_index is not None:
+            scores = scores.copy()
+            scores[exclude_index] = -np.inf
+        top = np.argpartition(-scores, k - 1)[:k] if k < n else np.arange(n)
+        top_indices: List[int] = top[np.argsort(-scores[top])].tolist()
+        return [self.examples[i] for i in top_indices]

pyproject.toml

@@ -54,6 +54,7 @@ jobs = [
     "dynaconf>=3.2.5,<4.0",
     "pydantic>=2.13.4,<3.0",
 ]
+llm = ["scikit-learn>=1.6,<2.0"]
 
 [build-system]
 requires = ["uv_build>=0.12.17,<0.13"]

tests/testutils/mock_example_retriever.py

@@ -0,0 +1,11 @@
+import numpy as np
+
+from machine.translation.llm.top_k_example_retriever import TopKExampleRetriever
+
+
+class MockExampleRetriever(TopKExampleRetriever):
+    def _fit(self) -> None:
+        pass
+
+    def _score_examples(self, query: str) -> np.ndarray:
+        return np.array([0 if query == e.source else (1 if query > e.source else -1) for e in self.examples])

tests/translation/llm/test_coverage_example_retriever.py

@@ -0,0 +1,76 @@
+from machine.translation.llm import Example
+from machine.translation.llm.coverage_example_retriever import CoverageExampleRetriever
+
+
+def test_coverage_retriever_prefers_an_uncovered_clause_to_a_near_duplicate():
+    sources = [
+        "the shepherd counted his sheep in the green field",
+        "the shepherd counted his sheep in the green field again",
+        "the shepherd counted all his sheep in the green field",
+        "fishermen mended nets",
+        "a woman drew water from the well",
+    ]
+    query = "the shepherd counted his sheep in the green field while fishermen mended nets"
+
+    retriever = CoverageExampleRetriever([Example(source, "") for source in sources])
+
+    assert [example.source for example in retriever.rank(query, 2)] == [sources[0], sources[3]]
+
+
+def test_coverage_retriever_keeps_the_closest_example():
+    sources = ["the shepherd counted his sheep in the green field", "fishermen mended nets"]
+    query = "the shepherd counted his sheep in the green field while fishermen mended nets"
+
+    retriever = CoverageExampleRetriever([Example(source, "") for source in sources])
+
+    assert [example.source for example in retriever.rank(query, 1)] == [sources[0]]
+
+
+def test_coverage_retriever_selects_an_exact_duplicate_pair_only_once():
+    examples = [Example("apple pie", "tarte"), Example("apple pie", "tarte"), Example("banana split", "banane")]
+    retriever = CoverageExampleRetriever(examples)
+
+    selected = retriever.rank("apple pie", 2)
+
+    assert examples[2] in selected
+    assert len([example for example in selected if example == examples[0]]) == 1
+
+
+def test_coverage_retriever_keeps_repeated_sources_with_different_targets():
+    examples = [Example("apple pie", "tarte"), Example("apple pie", "tourte"), Example("banana split", "banane")]
+
+    assert set(CoverageExampleRetriever(examples).rank("apple pie", 2)) == set(examples[:2])
+
+
+def test_coverage_retriever_excludes_copies_of_the_excluded_example():
+    examples = [Example("apple pie", "tarte"), Example("apple pie", "tarte"), Example("banana split", "banane")]
+
+    assert CoverageExampleRetriever(examples).rank("apple pie", 2, exclude=0) == [examples[2]]
+
+
+def test_coverage_retriever_finds_nothing_when_all_remaining_examples_are_copies():
+    examples = [Example("apple pie", "tarte"), Example("apple pie", "tarte")]
+
+    assert CoverageExampleRetriever(examples).rank("apple pie", 1, exclude=0) == []
+
+
+def test_coverage_retriever_matches_a_script_without_spaces():
+    sources = ["他们在家里吃饭", "我们去学校", "明天下雨"]
+
+    retriever = CoverageExampleRetriever([Example(source, "") for source in sources])
+
+    assert retriever.rank("我们明天去学校吗", 1)[0].source == sources[1]
+
+
+def test_coverage_retriever_fills_k_when_the_query_has_no_known_ngrams():
+    retriever = CoverageExampleRetriever(
+        [Example("apple pie", ""), Example("banana split", ""), Example("cherry tart", "")]
+    )
+
+    assert len(set(retriever.rank("zzzz", 2))) == 2
+
+
+def test_coverage_retriever_handles_a_corpus_with_no_text():
+    retriever = CoverageExampleRetriever([Example("", ""), Example("  ", "")])
+
+    assert retriever.rank("anything", 1) == []