diff --git a/src/SIL.Machine/Translation/Llm/CoverageExampleRetriever.cs b/src/SIL.Machine/Translation/Llm/CoverageExampleRetriever.cs new file mode 100644 index 000000000..9dd19b14a --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/CoverageExampleRetriever.cs @@ -0,0 +1,213 @@ +using System; +using System.Collections.Generic; +using System.Linq; + +namespace SIL.Machine.Translation.Llm +{ + public class CoverageExampleRetriever : ExampleRetriever + { + private const int Anchors = 20; + private const int CandidateMultiple = 4; + private const int MinCandidates = 100; + private const double RelevanceWeight = 0.7; + + private readonly HashSet _uniqueExampleIndices = new HashSet(); + private readonly HashSet _deduplicatedExampleIndices = new HashSet(); + private TfidfIndex _index; + + public CoverageExampleRetriever(IEnumerable examples) + : base(examples) + { + var seen = new Dictionary(); + for (int i = 0; i < ExampleList.Count; i++) + { + Example example = ExampleList[i]; + if (!seen.TryGetValue(example, out int firstIndex)) + { + _uniqueExampleIndices.Add(i); + _deduplicatedExampleIndices.Add(i); + seen[example] = i; + } + else + { + _uniqueExampleIndices.Remove(firstIndex); + } + } + Fit(); + } + + protected override void Fit() + { + _index = TfidfIndex.Create(ExampleList.Select(e => CharWbNgrams(e.Source)).ToArray()); + } + + protected override IReadOnlyList RankCore(string query, int k, int? excludeIndex) + { + if (_index == null) + return Array.Empty(); + + IReadOnlyDictionary queryVector = _index.Vectorize(CharWbNgrams(query)); + double[] similarities = _index.Similarities(queryVector); + // Highest similarity first; ties keep corpus order. + List ranked = Enumerable + .Range(0, similarities.Length) + .OrderByDescending(i => similarities[i]) + .ThenBy(i => i) + .ToList(); + if (k >= ranked.Count) + return ranked.Where(i => i != excludeIndex).Select(i => ExampleList[i]).ToList(); + + // If excludeIndex points to the first occurrence of a multiply attested example, + // that example is still included through its later copies. + if ( + excludeIndex.HasValue + && ( + !_deduplicatedExampleIndices.Contains(excludeIndex.Value) + || _uniqueExampleIndices.Contains(excludeIndex.Value) + ) + ) + { + ranked.RemoveAll(i => i == excludeIndex.Value); + } + ranked = ranked.Where(_deduplicatedExampleIndices.Contains).ToList(); + k = Math.Min(k, ranked.Count); + if (k == 0) + return Array.Empty(); + if (queryVector.Count == 0) + return ranked.Take(k).Select(i => ExampleList[i]).ToList(); + + int limit = Math.Max(CandidateMultiple * k, MinCandidates); + int[] candidates = ranked.Take(limit).ToArray(); + List chosen = ChooseForCoverage( + queryVector, + candidates.Select(i => _index.Rows[i]).ToArray(), + candidates.Select(i => similarities[i]).ToArray(), + k + ); + return chosen + .Select(i => candidates[i]) + .OrderByDescending(i => similarities[i]) + .ThenBy(i => i) + .Select(i => ExampleList[i]) + .ToList(); + } + + private static List ChooseForCoverage( + IReadOnlyDictionary queryVector, + IReadOnlyList> candidateRows, + double[] similarities, + int k + ) + { + // Restrict the features that contribute to query coverage to those present in the query. + int[] queryTerms = queryVector.Keys.ToArray(); + int candidateCount = candidateRows.Count; + var features = new bool[candidateCount][]; + for (int c = 0; c < candidateCount; c++) + { + features[c] = new bool[queryTerms.Length]; + for (int f = 0; f < queryTerms.Length; f++) + features[c][f] = candidateRows[c].TryGetValue(queryTerms[f], out double v) && v > 0; + } + + // Normalize the query weights and the candidate similarities. + double weightSum = queryTerms.Sum(t => queryVector[t]); + double[] weights = queryTerms.Select(t => queryVector[t] / weightSum).ToArray(); + double maxSimilarity = Math.Max(similarities.Max(), 1e-12); + double[] relevance = similarities.Select(s => s / maxSimilarity).ToArray(); + + // Always choose the first anchorCount most similar candidates. + int anchorCount = Math.Min(Anchors, Math.Min(Math.Max(1, k / 2), candidateCount)); + var chosen = Enumerable.Range(0, anchorCount).ToList(); + var available = new bool[candidateCount]; + for (int c = 0; c < candidateCount; c++) + available[c] = c >= anchorCount; + + // Track how many times each feature is represented across the chosen examples. + var counts = new double[queryTerms.Length]; + foreach (int c in chosen) + AddFeatures(counts, features[c]); + + // Greedily add the example with the highest composite relevance/coverage score. + int target = Math.Min(k, candidateCount); + var gains = new double[candidateCount]; + while (chosen.Count < target) + { + // The more important the n-gram is in the query, the higher the gain; the more times a chosen + // example already has it, the lower the gain. + var featureWeights = new double[queryTerms.Length]; + for (int f = 0; f < featureWeights.Length; f++) + featureWeights[f] = weights[f] / (1.0 + counts[f]); + double maxGain = double.NegativeInfinity; + for (int c = 0; c < candidateCount; c++) + { + double gain = 0; + for (int f = 0; f < featureWeights.Length; f++) + { + if (features[c][f]) + gain += featureWeights[f]; + } + gains[c] = gain; + if (available[c] && gain > maxGain) + maxGain = gain; + } + // Normalize the gains so they are comparable to the normalized relevance. + maxGain = Math.Max(maxGain, 1e-12); + + int best = -1; + double bestUtility = double.NegativeInfinity; + for (int c = 0; c < candidateCount; c++) + { + if (!available[c]) + continue; + double utility = RelevanceWeight * relevance[c] + (1 - RelevanceWeight) * (gains[c] / maxGain); + if (best < 0 || utility > bestUtility) + { + best = c; + bestUtility = utility; + } + } + chosen.Add(best); + available[best] = false; + AddFeatures(counts, features[best]); + } + return chosen; + } + + private static void AddFeatures(double[] counts, bool[] features) + { + for (int f = 0; f < features.Length; f++) + { + if (features[f]) + counts[f]++; + } + } + + /// Character 2- to 5-grams within word boundaries, as scikit-learn's char_wb does. + private static IReadOnlyList CharWbNgrams(string text) + { + var ngrams = new List(); + foreach (string word in text.ToLowerInvariant().Split((char[])null, StringSplitOptions.RemoveEmptyEntries)) + { + string padded = " " + word + " "; + for (int n = 2; n <= 5; n++) + { + int offset = 0; + ngrams.Add(Slice(padded, offset, n)); + while (offset + n < padded.Length) + { + offset++; + ngrams.Add(Slice(padded, offset, n)); + } + // A word shorter than n is counted only once. + if (offset == 0) + break; + } + } + return ngrams; + } + + private static string Slice(string text, int start, int length) => + text.Substring(start, Math.Min(length, text.Length - start)); + } +} diff --git a/src/SIL.Machine/Translation/Llm/Example.cs b/src/SIL.Machine/Translation/Llm/Example.cs new file mode 100644 index 000000000..7acdabd33 --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/Example.cs @@ -0,0 +1,33 @@ +using System; + +namespace SIL.Machine.Translation.Llm +{ + public sealed class Example : IEquatable + { + public Example(string source, string target) + { + Source = source ?? throw new ArgumentNullException(nameof(source)); + Target = target ?? throw new ArgumentNullException(nameof(target)); + } + + public string Source { get; } + public string Target { get; } + + public bool Equals(Example other) => + other != null + && string.Equals(Source, other.Source, StringComparison.Ordinal) + && string.Equals(Target, other.Target, StringComparison.Ordinal); + + public override bool Equals(object obj) => Equals(obj as Example); + + public override int GetHashCode() + { + unchecked + { + return (StringComparer.Ordinal.GetHashCode(Source) * 397) ^ StringComparer.Ordinal.GetHashCode(Target); + } + } + + public override string ToString() => $"Example(Source={Source}, Target={Target})"; + } +} diff --git a/src/SIL.Machine/Translation/Llm/ExampleRetrievalMethod.cs b/src/SIL.Machine/Translation/Llm/ExampleRetrievalMethod.cs new file mode 100644 index 000000000..eabcbf60b --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/ExampleRetrievalMethod.cs @@ -0,0 +1,8 @@ +namespace SIL.Machine.Translation.Llm +{ + public enum ExampleRetrievalMethod + { + Tfidf, + Coverage, + } +} diff --git a/src/SIL.Machine/Translation/Llm/ExampleRetriever.cs b/src/SIL.Machine/Translation/Llm/ExampleRetriever.cs new file mode 100644 index 000000000..db83d028b --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/ExampleRetriever.cs @@ -0,0 +1,46 @@ +using System; +using System.Collections.Generic; +using System.Linq; + +namespace SIL.Machine.Translation.Llm +{ + public abstract class ExampleRetriever + { + private readonly List _examples; + + protected ExampleRetriever(IEnumerable examples) + { + if (examples == null) + throw new ArgumentNullException(nameof(examples)); + _examples = examples.ToList(); + } + + /// A copy of the examples this retriever was built from. + public IReadOnlyList Examples => _examples.ToList(); + + protected IReadOnlyList ExampleList => _examples; + + /// + /// Builds the retrieval index. Each most-derived constructor must call this once its own state is set up. + /// + protected abstract void Fit(); + + /// + /// Returns up to examples ranked for . + /// names a position in the example list to leave out; an + /// out-of-range value excludes nothing. + /// + public IReadOnlyList Rank(string query, int k, int? excludeIndex = null) + { + if (query == null) + throw new ArgumentNullException(nameof(query)); + if (k <= 0) + return Array.Empty(); + if (excludeIndex.HasValue && (excludeIndex.Value < 0 || excludeIndex.Value >= _examples.Count)) + excludeIndex = null; + return RankCore(query, k, excludeIndex); + } + + protected abstract IReadOnlyList RankCore(string query, int k, int? excludeIndex); + } +} diff --git a/src/SIL.Machine/Translation/Llm/ExampleRetrieverFactory.cs b/src/SIL.Machine/Translation/Llm/ExampleRetrieverFactory.cs new file mode 100644 index 000000000..ad7f1263e --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/ExampleRetrieverFactory.cs @@ -0,0 +1,27 @@ +using System; +using System.Collections.Generic; + +namespace SIL.Machine.Translation.Llm +{ + public static class ExampleRetrieverFactory + { + public static ExampleRetriever CreateExampleRetriever( + ExampleRetrievalMethod method, + IEnumerable examples + ) + { + switch (method) + { + case ExampleRetrievalMethod.Tfidf: + return new TfidfExampleRetriever(examples); + case ExampleRetrievalMethod.Coverage: + return new CoverageExampleRetriever(examples); + default: + throw new ArgumentException( + $"Unable to create example retriever with method {method}.", + nameof(method) + ); + } + } + } +} diff --git a/src/SIL.Machine/Translation/Llm/TfidfExampleRetriever.cs b/src/SIL.Machine/Translation/Llm/TfidfExampleRetriever.cs new file mode 100644 index 000000000..c75843997 --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/TfidfExampleRetriever.cs @@ -0,0 +1,40 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using SIL.Machine.Tokenization; + +namespace SIL.Machine.Translation.Llm +{ + public class TfidfExampleRetriever : TopKExampleRetriever + { + private readonly LatinWordTokenizer _tokenizer = new LatinWordTokenizer(); + private TfidfIndex _index; + + public TfidfExampleRetriever(IEnumerable examples) + : base(examples) + { + Fit(); + } + + protected override void Fit() + { + _index = TfidfIndex.Create(ExampleList.Select(e => Tokenize(e.Source)).ToArray()); + } + + protected override double[] ScoreExamples(string query) + { + if (_index == null) + return Array.Empty(); + return _index.Similarities(_index.Vectorize(Tokenize(query))); + } + + private IReadOnlyList Tokenize(string text) + { + // Strip all punctuation-only tokens. + return _tokenizer + .Tokenize(text.ToLowerInvariant()) + .Where(token => token.Any(char.IsLetterOrDigit)) + .ToArray(); + } + } +} diff --git a/src/SIL.Machine/Translation/Llm/TfidfIndex.cs b/src/SIL.Machine/Translation/Llm/TfidfIndex.cs new file mode 100644 index 000000000..076a51739 --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/TfidfIndex.cs @@ -0,0 +1,91 @@ +using System; +using System.Collections.Generic; +using System.Linq; + +namespace SIL.Machine.Translation.Llm +{ + /// + /// A sparse, L2-normalized TF-IDF index over pre-tokenized documents. It follows scikit-learn's + /// TfidfVectorizer defaults: raw term counts, smoothed idf, unit-length rows. + /// + internal class TfidfIndex + { + private readonly Dictionary _vocabulary = new Dictionary(StringComparer.Ordinal); + private readonly List _idf = new List(); + + /// Returns null when no document has any term. + public static TfidfIndex Create(IReadOnlyList> documents) + { + var index = new TfidfIndex(documents); + return index._vocabulary.Count == 0 ? null : index; + } + + private TfidfIndex(IReadOnlyList> documents) + { + var documentFrequency = new List(); + foreach (IReadOnlyList document in documents) + { + foreach (string term in new HashSet(document, StringComparer.Ordinal)) + { + if (!_vocabulary.TryGetValue(term, out int termIndex)) + { + termIndex = _vocabulary.Count; + _vocabulary[term] = termIndex; + documentFrequency.Add(0); + } + documentFrequency[termIndex]++; + } + } + int n = documents.Count; + foreach (int df in documentFrequency) + _idf.Add(Math.Log((1.0 + n) / (1.0 + df)) + 1.0); + + Rows = documents.Select(Vectorize).ToArray(); + } + + /// The normalized vector of each indexed document, keyed by term index. + public IReadOnlyList> Rows { get; } + + /// Vectorizes a query. Terms outside the indexed vocabulary are dropped. + public IReadOnlyDictionary Vectorize(IReadOnlyList terms) + { + var vector = new Dictionary(); + foreach (string term in terms) + { + if (_vocabulary.TryGetValue(term, out int termIndex)) + { + vector.TryGetValue(termIndex, out double count); + vector[termIndex] = count + 1; + } + } + foreach (int termIndex in vector.Keys.ToList()) + vector[termIndex] *= _idf[termIndex]; + double norm = Math.Sqrt(vector.Values.Sum(v => v * v)); + if (norm > 0) + { + foreach (int termIndex in vector.Keys.ToList()) + vector[termIndex] /= norm; + } + return vector; + } + + /// Cosine similarity of the query vector to every indexed document. + public double[] Similarities(IReadOnlyDictionary query) + { + var similarities = new double[Rows.Count]; + if (query.Count == 0) + return similarities; + for (int i = 0; i < Rows.Count; i++) + { + double sum = 0; + foreach (KeyValuePair entry in query) + { + if (Rows[i].TryGetValue(entry.Key, out double value)) + sum += entry.Value * value; + } + similarities[i] = sum; + } + return similarities; + } + } +} diff --git a/src/SIL.Machine/Translation/Llm/TopKExampleRetriever.cs b/src/SIL.Machine/Translation/Llm/TopKExampleRetriever.cs new file mode 100644 index 000000000..78146a2c0 --- /dev/null +++ b/src/SIL.Machine/Translation/Llm/TopKExampleRetriever.cs @@ -0,0 +1,38 @@ +using System; +using System.Collections.Generic; +using System.Linq; + +namespace SIL.Machine.Translation.Llm +{ + public abstract class TopKExampleRetriever : ExampleRetriever + { + protected TopKExampleRetriever(IEnumerable examples) + : base(examples) { } + + /// One score per example, or an empty array when nothing is indexed. + protected abstract double[] ScoreExamples(string query); + + protected override IReadOnlyList RankCore(string query, int k, int? excludeIndex) + { + double[] scores = ScoreExamples(query); + if (scores.Length == 0) + return ExampleList.Where((_, i) => i != excludeIndex).Take(k).ToList(); + + int available = scores.Length - (excludeIndex.HasValue ? 1 : 0); + k = Math.Min(k, available); + if (k <= 0) + return Array.Empty(); + if (excludeIndex.HasValue) + scores[excludeIndex.Value] = double.NegativeInfinity; + + // Highest score first; ties keep corpus order. + return Enumerable + .Range(0, scores.Length) + .OrderByDescending(i => scores[i]) + .ThenBy(i => i) + .Take(k) + .Select(i => ExampleList[i]) + .ToList(); + } + } +} diff --git a/tests/SIL.Machine.Tests/Translation/Llm/CoverageExampleRetrieverTests.cs b/tests/SIL.Machine.Tests/Translation/Llm/CoverageExampleRetrieverTests.cs new file mode 100644 index 000000000..fb411ab5b --- /dev/null +++ b/tests/SIL.Machine.Tests/Translation/Llm/CoverageExampleRetrieverTests.cs @@ -0,0 +1,115 @@ +using NUnit.Framework; + +namespace SIL.Machine.Translation.Llm; + +[TestFixture] +public class CoverageExampleRetrieverTests +{ + private static CoverageExampleRetriever Create(params string[] sources) => + new CoverageExampleRetriever(sources.Select(s => new Example(s, ""))); + + [Test] + public void Rank_PrefersAnUncoveredClauseToANearDuplicate() + { + string[] 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", + }; + string query = "the shepherd counted his sheep in the green field while fishermen mended nets"; + + Assert.That(Create(sources).Rank(query, 2).Select(e => e.Source), Is.EqualTo(new[] { sources[0], sources[3] })); + } + + [Test] + public void Rank_KeepsTheClosestExample() + { + string[] sources = { "the shepherd counted his sheep in the green field", "fishermen mended nets" }; + string query = "the shepherd counted his sheep in the green field while fishermen mended nets"; + + Assert.That(Create(sources).Rank(query, 1).Select(e => e.Source), Is.EqualTo(new[] { sources[0] })); + } + + [Test] + public void Rank_SelectsAnExactDuplicatePairOnlyOnce() + { + var examples = new[] + { + new Example("apple pie", "tarte"), + new Example("apple pie", "tarte"), + new Example("banana split", "banane"), + }; + var retriever = new CoverageExampleRetriever(examples); + + var selected = retriever.Rank("apple pie", 2); + Assert.That(selected, Does.Contain(examples[2])); + Assert.That(selected.Count(e => e.Equals(examples[0])), Is.EqualTo(1)); + + selected = retriever.Rank("apple pie", 3); + Assert.That(selected.Count(e => e.Equals(examples[0])), Is.EqualTo(2)); + } + + [Test] + public void Rank_KeepsRepeatedSourcesWithDifferentTargets() + { + var examples = new[] + { + new Example("apple pie", "tarte"), + new Example("apple pie", "tourte"), + new Example("banana split", "banane"), + }; + + Assert.That(new CoverageExampleRetriever(examples).Rank("apple pie", 2), Is.EquivalentTo(examples.Take(2))); + } + + [Test] + public void Rank_KeepsACopyAtADifferentIndex() + { + var examples = new[] + { + new Example("apple pie", "tarte"), + new Example("apple pie", "tarte"), + new Example("banana split", "banane"), + }; + + Assert.That( + new CoverageExampleRetriever(examples).Rank("apple pie", 1, excludeIndex: 0), + Is.EqualTo(new[] { examples[1] }) + ); + } + + [Test] + public void Rank_ReturnsACopyWhenTheFirstPositionIsExcluded() + { + var examples = new[] { new Example("apple pie", "tarte"), new Example("apple pie", "tarte") }; + + Assert.That( + new CoverageExampleRetriever(examples).Rank("apple pie", 1, excludeIndex: 0), + Is.EqualTo(new[] { examples[1] }) + ); + } + + [Test] + public void Rank_MatchesAScriptWithoutSpaces() + { + string[] sources = { "他们在家里吃饭", "我们去学校", "明天下雨" }; + + Assert.That(Create(sources).Rank("我们明天去学校吗", 1)[0].Source, Is.EqualTo(sources[1])); + } + + [Test] + public void Rank_FillsKWhenTheQueryHasNoKnownNgrams() + { + var retriever = Create("apple pie", "banana split", "cherry tart"); + Assert.That(retriever.Rank("zzzz", 2).Distinct().Count(), Is.EqualTo(2)); + } + + [Test] + public void Rank_CorpusWithNoText() + { + Assert.That(Create("", " ").Rank("anything", 1), Is.Empty); + } +} diff --git a/tests/SIL.Machine.Tests/Translation/Llm/ExampleRetrieverFactoryTests.cs b/tests/SIL.Machine.Tests/Translation/Llm/ExampleRetrieverFactoryTests.cs new file mode 100644 index 000000000..f69efd7e5 --- /dev/null +++ b/tests/SIL.Machine.Tests/Translation/Llm/ExampleRetrieverFactoryTests.cs @@ -0,0 +1,33 @@ +using NUnit.Framework; + +namespace SIL.Machine.Translation.Llm; + +[TestFixture] +public class ExampleRetrieverFactoryTests +{ + [Test] + public void CreateExampleRetriever_Tfidf() + { + var examples = new[] { new Example("a", "0"), new Example("b", "1") }; + var retriever = ExampleRetrieverFactory.CreateExampleRetriever(ExampleRetrievalMethod.Tfidf, examples); + Assert.That(retriever, Is.TypeOf()); + Assert.That(retriever.Examples, Is.EqualTo(examples)); + } + + [Test] + public void CreateExampleRetriever_Coverage() + { + var examples = new[] { new Example("a", "0"), new Example("b", "1") }; + var retriever = ExampleRetrieverFactory.CreateExampleRetriever(ExampleRetrievalMethod.Coverage, examples); + Assert.That(retriever, Is.TypeOf()); + Assert.That(retriever.Examples, Is.EqualTo(examples)); + } + + [Test] + public void CreateExampleRetriever_UnknownMethod() + { + Assert.Throws(() => + ExampleRetrieverFactory.CreateExampleRetriever((ExampleRetrievalMethod)99, Array.Empty()) + ); + } +} diff --git a/tests/SIL.Machine.Tests/Translation/Llm/ExampleRetrieverTests.cs b/tests/SIL.Machine.Tests/Translation/Llm/ExampleRetrieverTests.cs new file mode 100644 index 000000000..572495d2c --- /dev/null +++ b/tests/SIL.Machine.Tests/Translation/Llm/ExampleRetrieverTests.cs @@ -0,0 +1,50 @@ +using NUnit.Framework; + +namespace SIL.Machine.Translation.Llm; + +[TestFixture] +public class ExampleRetrieverTests +{ + [Test] + public void Rank_KGreaterThanN() + { + var examples = new[] { new Example("a", "1"), new Example("b", "2") }; + var retriever = new MockExampleRetriever(examples); + Assert.That(retriever.Rank("query", 3), Is.EqualTo(examples)); + } + + [Test] + public void Rank_NoSuchExclude() + { + var examples = new[] { new Example("a", "1"), new Example("b", "2") }; + var retriever = new MockExampleRetriever(examples); + Assert.That(retriever.Rank("query", 2, excludeIndex: 10), Is.EqualTo(examples)); + } + + [Test] + public void Rank_ExcludesOnlyTheSpecifiedIndex() + { + var duplicate = new Example("a", "1"); + var examples = new[] { duplicate, new Example("b", "2"), duplicate }; + var retriever = new MockExampleRetriever(examples); + Assert.That(retriever.Rank("query", 3, excludeIndex: 0), Is.EqualTo(examples.Skip(1))); + } + + [Test] + public void Rank_NonPositiveK() + { + var retriever = new MockExampleRetriever(new[] { new Example("a", "1") }); + Assert.That(retriever.Rank("query", 0), Is.Empty); + } + + [Test] + public void Examples_ReturnsACopy() + { + var examples = new[] { new Example("a", "1"), new Example("b", "2") }; + var retriever = new MockExampleRetriever(examples); + + ((List)retriever.Examples).Clear(); + + Assert.That(retriever.Examples, Is.EqualTo(examples)); + } +} diff --git a/tests/SIL.Machine.Tests/Translation/Llm/MockExampleRetriever.cs b/tests/SIL.Machine.Tests/Translation/Llm/MockExampleRetriever.cs new file mode 100644 index 000000000..4c17b5e20 --- /dev/null +++ b/tests/SIL.Machine.Tests/Translation/Llm/MockExampleRetriever.cs @@ -0,0 +1,25 @@ +namespace SIL.Machine.Translation.Llm; + +public class MockExampleRetriever : TopKExampleRetriever +{ + public MockExampleRetriever(IEnumerable examples) + : base(examples) + { + Fit(); + } + + protected override void Fit() { } + + protected override double[] ScoreExamples(string query) + { + return ExampleList + .Select(e => + { + int comparison = string.CompareOrdinal(query, e.Source); + return comparison == 0 ? 0.0 + : comparison > 0 ? 1.0 + : -1.0; + }) + .ToArray(); + } +} diff --git a/tests/SIL.Machine.Tests/Translation/Llm/TfidfExampleRetrieverTests.cs b/tests/SIL.Machine.Tests/Translation/Llm/TfidfExampleRetrieverTests.cs new file mode 100644 index 000000000..bc13b8957 --- /dev/null +++ b/tests/SIL.Machine.Tests/Translation/Llm/TfidfExampleRetrieverTests.cs @@ -0,0 +1,99 @@ +using NUnit.Framework; + +namespace SIL.Machine.Translation.Llm; + +[TestFixture] +public class TfidfExampleRetrieverTests +{ + [Test] + public void Rank_MostSimilarSourceFirst() + { + var examples = new[] + { + new Example("the cat sat on the mat", "0"), + new Example("completely unrelated sentence", "1"), + new Example("a cat sat here", "2"), + }; + + Assert.That( + new TfidfExampleRetriever(examples).Rank("the cat sat", 2), + Is.EqualTo(new[] { examples[0], examples[2] }) + ); + } + + [Test] + public void Rank_RespectsK() + { + var examples = Enumerable.Range(0, 3).Select(i => new Example("apple pie", i.ToString())).ToArray(); + Assert.That(new TfidfExampleRetriever(examples).Rank("apple pie", 2), Has.Count.EqualTo(2)); + } + + [Test] + public void Rank_ExcludesItsOwnPosition() + { + var examples = new[] + { + new Example("the cat sat", "0"), + new Example("the cat sat", "1"), + new Example("totally different words", "2"), + }; + + Assert.That( + new TfidfExampleRetriever(examples).Rank("the cat sat", 2, excludeIndex: 0), + Is.EqualTo(new[] { examples[1], examples[2] }) + ); + } + + [Test] + public void Rank_KeepsDuplicateCopiesAtOtherIndices() + { + var excluded = new Example("apple pie", "tarte"); + var retriever = new TfidfExampleRetriever( + new[] { excluded, excluded, new Example("banana split", "banane"), new Example("cherry tart", "cerise") } + ); + + Assert.That(retriever.Rank("apple pie", 1, excludeIndex: 0), Is.EqualTo(new[] { excluded })); + } + + [Test] + public void Rank_KLargerThanTheCorpusReturnsEveryExample() + { + var examples = new[] { new Example("apple pie", "0"), new Example("banana split", "1") }; + Assert.That(new TfidfExampleRetriever(examples).Rank("apple pie", 10), Is.EqualTo(examples)); + } + + [Test] + public void Rank_EmptyCorpus() + { + Assert.That(new TfidfExampleRetriever(Array.Empty()).Rank("anything", 5), Is.Empty); + } + + [Test] + public void Rank_CorpusWithNoWordTokens() + { + var examples = new[] { new Example("!!!", "0"), new Example("...", "1") }; + + Assert.That(new TfidfExampleRetriever(examples).Rank("anything", 1), Is.EqualTo(new[] { examples[0] })); + Assert.That(new TfidfExampleRetriever(examples).Rank("anything", 4), Is.EqualTo(examples)); + } + + [Test] + public void Rank_OnlySomeSourcesHaveWords() + { + var examples = new[] { new Example("!!!", "0"), new Example("let there be light", "1") }; + Assert.That(new TfidfExampleRetriever(examples).Rank("light", 1), Is.EqualTo(new[] { examples[1] })); + } + + [Test] + public void Rank_PunctuationCarriesNoSignal() + { + var examples = new[] + { + new Example("the cat sat", "0"), + new Example("!!! !!! !!! ???", "1"), + new Example("something else entirely", "2"), + }; + + Assert.That(new TfidfExampleRetriever(examples).Rank("!!! !!! !!! cat", 1), Is.EqualTo(new[] { examples[0] })); + } +}