Port 'LLM example retrieval' from machine.py

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<int> _uniqueExampleIndices = new HashSet<int>();
+        private readonly HashSet<int> _deduplicatedExampleIndices = new HashSet<int>();
+        private TfidfIndex _index;
+
+        public CoverageExampleRetriever(IEnumerable<Example> examples)
+            : base(examples)
+        {
+            var seen = new Dictionary<Example, int>();
+            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<Example> RankCore(string query, int k, int? excludeIndex)
+        {
+            if (_index == null)
+                return Array.Empty<Example>();
+
+            IReadOnlyDictionary<int, double> queryVector = _index.Vectorize(CharWbNgrams(query));
+            double[] similarities = _index.Similarities(queryVector);
+            // Highest similarity first; ties keep corpus order.
+            List<int> 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<Example>();
+            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<int> 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<int> ChooseForCoverage(
+            IReadOnlyDictionary<int, double> queryVector,
+            IReadOnlyList<IReadOnlyDictionary<int, double>> 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]++;
+            }
+        }
+
+        /// <summary>Character 2- to 5-grams within word boundaries, as scikit-learn's char_wb does.</summary>
+        private static IReadOnlyList<string> CharWbNgrams(string text)
+        {
+            var ngrams = new List<string>();
+            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));
+    }
+}

src/SIL.Machine/Translation/Llm/Example.cs

@@ -0,0 +1,33 @@
+using System;
+
+namespace SIL.Machine.Translation.Llm
+{
+    public sealed class Example : IEquatable<Example>
+    {
+        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})";
+    }
+}

src/SIL.Machine/Translation/Llm/ExampleRetrievalMethod.cs

@@ -0,0 +1,8 @@
+namespace SIL.Machine.Translation.Llm
+{
+    public enum ExampleRetrievalMethod
+    {
+        Tfidf,
+        Coverage,
+    }
+}

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<Example> _examples;
+
+        protected ExampleRetriever(IEnumerable<Example> examples)
+        {
+            if (examples == null)
+                throw new ArgumentNullException(nameof(examples));
+            _examples = examples.ToList();
+        }
+
+        /// <summary>A copy of the examples this retriever was built from.</summary>
+        public IReadOnlyList<Example> Examples => _examples.ToList();
+
+        protected IReadOnlyList<Example> ExampleList => _examples;
+
+        /// <summary>
+        /// Builds the retrieval index. Each most-derived constructor must call this once its own state is set up.
+        /// </summary>
+        protected abstract void Fit();
+
+        /// <summary>
+        /// Returns up to <paramref name="k"/> examples ranked for <paramref name="query"/>.
+        /// <paramref name="excludeIndex"/> names a position in the example list to leave out; an
+        /// out-of-range value excludes nothing.
+        /// </summary>
+        public IReadOnlyList<Example> Rank(string query, int k, int? excludeIndex = null)
+        {
+            if (query == null)
+                throw new ArgumentNullException(nameof(query));
+            if (k <= 0)
+                return Array.Empty<Example>();
+            if (excludeIndex.HasValue && (excludeIndex.Value < 0 || excludeIndex.Value >= _examples.Count))
+                excludeIndex = null;
+            return RankCore(query, k, excludeIndex);
+        }
+
+        protected abstract IReadOnlyList<Example> RankCore(string query, int k, int? excludeIndex);
+    }
+}

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<Example> 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)
+                    );
+            }
+        }
+    }
+}

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<Example> 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<double>();
+            return _index.Similarities(_index.Vectorize(Tokenize(query)));
+        }
+
+        private IReadOnlyList<string> Tokenize(string text)
+        {
+            // Strip all punctuation-only tokens.
+            return _tokenizer
+                .Tokenize(text.ToLowerInvariant())
+                .Where(token => token.Any(char.IsLetterOrDigit))
+                .ToArray();
+        }
+    }
+}