LLM example retrieval

machine/translation/llm/__init__.py

@@ -0,0 +1,22 @@
+from ...utils.packages import is_sklearn_available
+
+if not is_sklearn_available():
+    raise RuntimeError('sklearn is not installed. Install sil-machine with the "llm" extra.')
+
+from .coverage_example_retriever import CoverageExampleRetriever
+from .example import Example
+from .example_retriever import ExampleRetriever
+from .example_retriever_factory import ExampleRetrieverFactory
+from .example_retriever_method import ExampleRetrievalMethod
+from .tfidf_example_retriever import TfidfExampleRetriever
+from .top_k_example_retriever import TopKExampleRetriever
+
+__all__ = [
+    "CoverageExampleRetriever",
+    "Example",
+    "ExampleRetrieverFactory",
+    "ExampleRetrievalMethod",
+    "ExampleRetriever",
+    "TfidfExampleRetriever",
+    "TopKExampleRetriever",
+]

machine/translation/llm/coverage_example_retriever.py

@@ -0,0 +1,117 @@
+from typing import 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 .example_retriever import ExampleRetriever
+
+
+class CoverageExampleRetriever(ExampleRetriever):
+    _ANCHORS = 20
+    _CANDIDATE_MULTIPLE = 4
+    _MIN_CANDIDATES = 100
+    _RELEVANCE_WEIGHT = 0.7
+
+    def __init__(self, examples: List[Example]) -> None:
+        super().__init__(examples)
+        seen = {}
+        self._unique_example_indices = set()
+        self._deduplicated_example_indices = []
+        for index, example in enumerate(self._examples):
+            if example not in seen:
+                self._unique_example_indices.add(index)
+                self._deduplicated_example_indices.append(index)
+                seen[example] = index
+            else:
+                if seen[example] in self._unique_example_indices:
+                    self._unique_example_indices.remove(seen[example])
+
+    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 []
+        query_vector = cast(csr_matrix, self._vectorizer.transform([query]))
+        similarities = linear_kernel(query_vector, self._matrix).ravel()
+        ranked = similarities.argsort()[::-1]
+        if k >= ranked.shape[0]:
+            return [self._examples[i] for i in ranked if i != exclude_index]
+        # If exclude index points to an example which is the first occurrence
+        # of a multiply attested example, then we still include that example
+        if exclude_index is not None and (
+            exclude_index not in self._deduplicated_example_indices or exclude_index in self._unique_example_indices
+        ):
+            ranked = ranked[ranked != exclude_index]
+        ranked = ranked[np.isin(ranked, self._deduplicated_example_indices)]
+        k = min(k, len(ranked))
+        if k == 0:
+            return []
+        if query_vector.nnz == 0:
+            return [self._examples[i] for i in ranked[:k].tolist()]
+        limit = max(self._CANDIDATE_MULTIPLE * k, self._MIN_CANDIDATES)
+        candidates = ranked[:limit]
+        if len(candidates) == 0:
+            return []
+        chosen = self._choose_for_coverage(query_vector, self._matrix[candidates], similarities[candidates], k)
+        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 _choose_for_coverage(
+        self, query_vector: csr_matrix, candidate_matrix: csr_matrix, similarities: np.ndarray, k: int
+    ) -> List[int]:
+        # Restrict the features that contribute to query coverage to those present in the query.
+        features = candidate_matrix[:, query_vector.indices].toarray() > 0
+
+        # Normalize the query weights and the candidate similarities.
+        weights = query_vector.data / query_vector.data.sum()
+        relevance = similarities / max(float(similarities.max()), 1e-12)
+
+        # Always choose the first anchor_count most similar candidates.
+        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
+
+        # Keep track of how many times each feature is represented across the chosen examples.
+        counts = features[chosen].sum(axis=0).astype(float)
+
+        # Greedily select additional examples by repeatedly selecting
+        # the example with the highest composite relevance/coverage score
+        while len(chosen) < min(k, len(features)):
+            # The more the n-gram feature is important in the query, the higher the gain;
+            # the more times we've already selected an example that has that feature, the lower the gain.
+            relevance_weighted_coverage_gain = features @ (weights / (1.0 + counts))
+
+            # Normalize the gains so they are comparable to the normalized relevance.
+            relevance_weighted_coverage_gain /= max(float(relevance_weighted_coverage_gain[available].max()), 1e-12)
+
+            # Create the composite utility score: a weighted sum of the gain and the normalized relevance score.
+            utility = (
+                self._RELEVANCE_WEIGHT * relevance + (1 - self._RELEVANCE_WEIGHT) * relevance_weighted_coverage_gain
+            )
+
+            # If the example has already been selected, set the score to -inf.
+            utility[~available] = -np.inf
+
+            # Greedily choose the highest-scoring candidate and update the available values and counts.
+            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.py

@@ -0,0 +1,29 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from .example import Example
+
+
+class ExampleRetriever(ABC):
+
+    def __init__(self, examples: List[Example]) -> None:
+        self._examples = list(examples)
+        self._fit()
+
+    @abstractmethod
+    def _fit(self) -> None: ...
+
+    @property
+    def examples(self) -> List[Example]:
+        return list(self._examples)
+
+    def rank(self, query: str, k: int, exclude_index: Optional[int] = None) -> List[Example]:
+        if k <= 0:
+            return []
+        num_examples = len(self._examples)
+        if exclude_index is not None and not 0 <= exclude_index < num_examples:
+            exclude_index = None
+        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/example_retriever_factory.py

@@ -0,0 +1,19 @@
+from typing import List
+
+from .coverage_example_retriever import CoverageExampleRetriever
+from .example import Example
+from .example_retriever import ExampleRetriever
+from .example_retriever_method import ExampleRetrievalMethod
+from .tfidf_example_retriever import TfidfExampleRetriever
+
+
+class ExampleRetrieverFactory:
+    @staticmethod
+    def create_example_retriever(method: ExampleRetrievalMethod, examples: List[Example]) -> ExampleRetriever:
+        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/tfidf_example_retriever.py

@@ -0,0 +1,38 @@
+from typing import List, 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: List[Example]) -> None:
+        self._latin_tokenizer = LatinWordTokenizer()
+        super().__init__(examples)
+
+    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([])
+        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,29 @@
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+import numpy as np
+
+from machine.translation.llm.example import Example
+
+from .example_retriever import ExampleRetriever
+
+
+class TopKExampleRetriever(ExampleRetriever, 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)
+        if len(scores) == 0:
+            return [example for i, example in enumerate(self._examples) if i != exclude_index][:k]
+        n = scores.shape[0]
+        available = n - (exclude_index is not None)
+        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]

machine/utils/packages.py

@@ -23,3 +23,7 @@ def is_tensorflow_available() -> bool:
 
 def is_sentencepiece_available() -> bool:
     return find_spec("sentencepiece") is not None
+
+
+def is_sklearn_available() -> bool:
+    return find_spec("sklearn") is not None

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.0 if query > e.source else -1.0) for e in self._examples])

tests/translation/llm/test_coverage_example_retriever.py

@@ -0,0 +1,79 @@
+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
+
+    selected = retriever.rank("apple pie", 3)
+    assert len([example for example in selected if example == examples[0]]) == 2
+
+
+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_keeps_a_copy_at_a_different_index():
+    examples = [Example("apple pie", "tarte"), Example("apple pie", "tarte"), Example("banana split", "banane")]
+
+    assert CoverageExampleRetriever(examples).rank("apple pie", 1, exclude_index=0) == [examples[1]]
+
+
+def test_coverage_retriever_returns_a_copy_when_the_first_position_is_excluded():
+    examples = [Example("apple pie", "tarte"), Example("apple pie", "tarte")]
+
+    assert CoverageExampleRetriever(examples).rank("apple pie", 1, exclude_index=0) == [examples[1]]
+
+
+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) == []