Port transductive alignment model

src/SIL.Machine.Translation.Thot/IParallelTextCorpusExtensions.cs

@@ -0,0 +1,130 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using SIL.Machine.Corpora;
+using SIL.Machine.Utils;
+
+namespace SIL.Machine.Translation.Thot
+{
+    public static class IParallelTextCorpusExtensions
+    {
+        public static IParallelTextCorpus WordAlign(
+            this IParallelTextCorpus corpus,
+            ThotWordAlignmentModelType modelType = ThotWordAlignmentModelType.FastAlign,
+            SymmetrizationHeuristic symmetrizationHeuristic = SymmetrizationHeuristic.GrowDiagFinalAnd,
+            IProgress<ProgressStatus> progress = null
+        ) => new TrainedWordAlignParallelTextCorpus(corpus, modelType, symmetrizationHeuristic, progress);
+
+        public static IParallelTextCorpus WordAlign(
+            this IParallelTextCorpus corpus,
+            ThotSymmetrizedWordAlignmentModel model,
+            int batchSize = 1024
+        )
+        {
+            if (model.EmitTrainingAlignments)
+                return new TransductiveWordAlignParallelTextCorpus(corpus, model);
+
+            return CorporaExtensions.WordAlign(corpus, model, batchSize);
+        }
+
+        private class TransductiveWordAlignParallelTextCorpus : WordAlignParallelTextCorpusBase
+        {
+            private readonly IParallelTextCorpus _corpus;
+            private readonly ITransductiveWordAlignmentModel _model;
+
+            public TransductiveWordAlignParallelTextCorpus(
+                IParallelTextCorpus corpus,
+                ITransductiveWordAlignmentModel model
+            )
+                : base(corpus)
+            {
+                _corpus = corpus;
+                _model = model;
+            }
+
+            public override IEnumerable<ParallelTextRow> GetRows(IEnumerable<string> textIds) =>
+                GetTransductiveRows(_corpus, _model, textIds);
+        }
+
+        private class TrainedWordAlignParallelTextCorpus : WordAlignParallelTextCorpusBase
+        {
+            private readonly IParallelTextCorpus _corpus;
+            private readonly ThotWordAlignmentModelType _modelType;
+            private readonly SymmetrizationHeuristic _symmetrizationHeuristic;
+            private readonly IProgress<ProgressStatus> _progress;
+
+            public TrainedWordAlignParallelTextCorpus(
+                IParallelTextCorpus corpus,
+                ThotWordAlignmentModelType modelType,
+                SymmetrizationHeuristic symmetrizationHeuristic,
+                IProgress<ProgressStatus> progress
+            )
+                : base(corpus)
+            {
+                _corpus = corpus;
+                _modelType = modelType;
+                _symmetrizationHeuristic = symmetrizationHeuristic;
+                _progress = progress;
+            }
+
+            public override IEnumerable<ParallelTextRow> GetRows(IEnumerable<string> textIds)
+            {
+                // Training on only the requested texts keeps the training-alignment index in sync with the rows.
+                IParallelTextCorpus corpus = _corpus.FilterTexts(textIds);
+                // Training in the generator ties the model's lifetime to reading the rows, at the cost of
+                // training a new model on each iteration.
+                using (var model = ThotSymmetrizedWordAlignmentModel.Create(_modelType))
+                {
+                    model.Heuristic = _symmetrizationHeuristic;
+                    // Retain the alignments computed during training so that the corpus can be aligned
+                    // without a separate, potentially expensive, inference pass.
+                    model.EmitTrainingAlignments = true;
+                    using (ITrainer trainer = model.CreateTrainer(corpus))
+                    {
+                        trainer.TrainAsync(_progress).GetAwaiter().GetResult();
+                        trainer.SaveAsync().GetAwaiter().GetResult();
+                    }
+
+                    foreach (ParallelTextRow row in GetTransductiveRows(corpus, model, textIds: null))
+                        yield return row;
+                }
+            }
+        }
+
+        private static IEnumerable<ParallelTextRow> GetTransductiveRows(
+            IParallelTextCorpus corpus,
+            ITransductiveWordAlignmentModel model,
+            IEnumerable<string> textIds
+        )
+        {
+            // The training alignments are keyed by the order in which the sentence pairs were added during
+            // training, so the corpus the model was trained on must be iterated in full to keep the index in
+            // sync; rows outside the requested texts are skipped rather than filtered out.
+            var textIdList = textIds?.ToList();
+            List<ParallelTextRow> rows = corpus.GetRows().ToList();
+            for (int i = 0; i < rows.Count; i++)
+            {
+                ParallelTextRow row = rows[i];
+                if (textIdList != null && !textIdList.Contains(row.TextId))
+                    continue;
+
+                WordAlignmentMatrix alignment = model.GetTrainingAlignment(i);
+                WordAlignmentMatrix knownAlignment = row.CreateAlignmentMatrix();
+                if (knownAlignment != null)
+                {
+                    knownAlignment.PrioritySymmetrizeWith(alignment);
+                    alignment = knownAlignment;
+                }
+
+                IReadOnlyCollection<AlignedWordPair> wordPairs = alignment.ToAlignedWordPairs();
+                if (model is IWordAlignmentModel wordAlignmentModel)
+                {
+                    wordAlignmentModel.ComputeAlignedWordPairScores(row.SourceSegment, row.TargetSegment, wordPairs);
+                }
+
+                row.AlignedWordPairs = wordPairs;
+                yield return row;
+            }
+        }
+    }
+}

src/SIL.Machine.Translation.Thot/SIL.Machine.Translation.Thot.csproj

@@ -12,7 +12,7 @@
 	<Import Project="../AssemblyInfo.props" />
 
 	<ItemGroup>
-		<PackageReference Include="Thot" Version="3.5.0" />
+		<PackageReference Include="Thot" Version="3.5.1" />
 	</ItemGroup>
 
 	<ItemGroup>

src/SIL.Machine.Translation.Thot/Thot.cs

@@ -172,6 +172,21 @@ uint capacity
         [DllImport("thot", CallingConvention = CallingConvention.Cdecl)]
         public static extern uint swAlignModel_getMaxSentenceLength(IntPtr swAlignModelHandle);
 
+        [DllImport("thot", CallingConvention = CallingConvention.Cdecl)]
+        public static extern uint swAlignModel_getNumSentencePairs(IntPtr swAlignModelHandle);
+
+        [DllImport("thot", CallingConvention = CallingConvention.Cdecl)]
+        public static extern double swAlignModel_getTrainingAlignment(
+            IntPtr swAlignModelHandle,
+            uint n,
+            IntPtr matrix,
+            ref uint iLen,
+            ref uint jLen
+        );
+
+        [DllImport("thot", CallingConvention = CallingConvention.Cdecl)]
+        public static extern void swAlignModel_setEmitTrainingAlignments(IntPtr swAlignModelHandle, bool value);
+
         [DllImport("thot", CallingConvention = CallingConvention.Cdecl)]
         public static extern void swAlignModel_setVariationalBayes(IntPtr swAlignModelHandle, bool variationalBayes);
 

src/SIL.Machine.Translation.Thot/ThotSymmetrizedWordAlignmentModel.cs

@@ -0,0 +1,62 @@
+namespace SIL.Machine.Translation.Thot
+{
+    public class ThotSymmetrizedWordAlignmentModel : SymmetrizedWordAlignmentModel, ITransductiveWordAlignmentModel
+    {
+        private readonly ThotWordAlignmentModel _directWordAlignmentModel;
+        private readonly ThotWordAlignmentModel _inverseWordAlignmentModel;
+
+        public ThotSymmetrizedWordAlignmentModel(
+            ThotWordAlignmentModel directWordAlignmentModel,
+            ThotWordAlignmentModel inverseWordAlignmentModel
+        )
+            : base(directWordAlignmentModel, inverseWordAlignmentModel)
+        {
+            _directWordAlignmentModel = directWordAlignmentModel;
+            _inverseWordAlignmentModel = inverseWordAlignmentModel;
+        }
+
+        public bool EmitTrainingAlignments
+        {
+            get => _directWordAlignmentModel.EmitTrainingAlignments;
+            set
+            {
+                _directWordAlignmentModel.EmitTrainingAlignments = value;
+                _inverseWordAlignmentModel.EmitTrainingAlignments = value;
+            }
+        }
+
+        public int TrainingAlignmentCount => _directWordAlignmentModel.TrainingAlignmentCount;
+
+        public static ThotSymmetrizedWordAlignmentModel Create(ThotWordAlignmentModelType modelType) =>
+            new ThotSymmetrizedWordAlignmentModel(
+                ThotWordAlignmentModel.Create(modelType),
+                ThotWordAlignmentModel.Create(modelType)
+            );
+
+        public WordAlignmentMatrix GetTrainingAlignment(int n)
+        {
+            WordAlignmentMatrix bestMatrix = _directWordAlignmentModel.GetTrainingAlignment(n);
+            if (Heuristic == SymmetrizationHeuristic.None)
+                return bestMatrix;
+
+            WordAlignmentMatrix invMatrix = _inverseWordAlignmentModel.GetTrainingAlignment(n);
+            invMatrix.Transpose();
+
+            // Skip the combine when the matrices are degenerate or their dimensions don't
+            // line up (e.g. an out-of-range n, or a pair filtered out of training in only
+            // one direction): the heuristic operations require matching dimensions.
+            if (
+                bestMatrix.RowCount == 0
+                || bestMatrix.ColumnCount == 0
+                || invMatrix.RowCount != bestMatrix.RowCount
+                || invMatrix.ColumnCount != bestMatrix.ColumnCount
+            )
+            {
+                return bestMatrix;
+            }
+
+            bestMatrix.SymmetrizeWith(invMatrix, Heuristic);
+            return bestMatrix;
+        }
+    }
+}

src/SIL.Machine.Translation.Thot/ThotWordAlignmentModel.cs

@@ -13,7 +13,10 @@
 
 namespace SIL.Machine.Translation.Thot
 {
-    public abstract class ThotWordAlignmentModel : DisposableBase, IIbm1WordAlignmentModel
+    public abstract class ThotWordAlignmentModel
+        : DisposableBase,
+            ITransductiveWordAlignmentModel,
+            IIbm1WordAlignmentModel
     {
         public static ThotWordAlignmentModel Create(ThotWordAlignmentModelType type)
         {
@@ -156,6 +159,30 @@ public void Save()
                 Thot.swAlignModel_save(Handle, _prefFileName);
         }
 
+        public bool EmitTrainingAlignments { get; set; }
+
+        public int TrainingAlignmentCount => (int)Thot.swAlignModel_getNumSentencePairs(Handle);
+
+        public WordAlignmentMatrix GetTrainingAlignment(int n)
+        {
+            CheckDisposed();
+
+            uint iLen = 0;
+            uint jLen = 0;
+            Thot.swAlignModel_getTrainingAlignment(Handle, (uint)n, IntPtr.Zero, ref iLen, ref jLen);
+
+            IntPtr nativeMatrix = Thot.AllocNativeMatrix((int)iLen, (int)jLen);
+            try
+            {
+                Thot.swAlignModel_getTrainingAlignment(Handle, (uint)n, nativeMatrix, ref iLen, ref jLen);
+                return Thot.ConvertNativeMatrixToWordAlignmentMatrix(nativeMatrix, iLen, jLen);
+            }
+            finally
+            {
+                Thot.FreeNativeMatrix(nativeMatrix, iLen);
+            }
+        }
+
         public double GetTranslationScore(string sourceWord, string targetWord)
         {
             return GetTranslationProbability(sourceWord, targetWord);
@@ -316,7 +343,7 @@ private class Trainer : ThotWordAlignmentModelTrainer
             private readonly ThotWordAlignmentModel _model;
 
             public Trainer(ThotWordAlignmentModel model, IParallelTextCorpus corpus)
-                : base(model.Type, corpus, model._prefFileName, model.Parameters)
+                : base(model.Type, corpus, model._prefFileName, model.Parameters, model.EmitTrainingAlignments)
             {
                 _model = model;
                 CloseOnDispose = false;

src/SIL.Machine.Translation.Thot/ThotWordAlignmentModelTrainer.cs

@@ -26,9 +26,10 @@ public ThotWordAlignmentModelTrainer(
             string sourceFileName,
             string targetFileName,
             string prefFileName,
-            ThotWordAlignmentParameters parameters = null
+            ThotWordAlignmentParameters parameters = null,
+            bool emitTrainingAlignments = false
         )
-            : this(modelType, null, prefFileName, parameters)
+            : this(modelType, null, prefFileName, parameters, emitTrainingAlignments)
         {
             _sourceFileName = sourceFileName;
             _targetFileName = targetFileName;
@@ -38,7 +39,8 @@ public ThotWordAlignmentModelTrainer(
             ThotWordAlignmentModelType modelType,
             IParallelTextCorpus corpus,
             string prefFileName,
-            ThotWordAlignmentParameters parameters = null
+            ThotWordAlignmentParameters parameters = null,
+            bool emitTrainingAlignments = false
         )
         {
             _prefFileName = prefFileName;
@@ -47,6 +49,8 @@ public ThotWordAlignmentModelTrainer(
             if (parameters == null)
                 parameters = new ThotWordAlignmentParameters();
 
+            EmitTrainingAlignments = emitTrainingAlignments;
+
             _models = new List<(IntPtr, int)>();
             if (modelType == ThotWordAlignmentModelType.FastAlign)
             {
@@ -197,6 +201,8 @@ public ThotWordAlignmentModelTrainer(
 
         public TrainStats Stats { get; } = new TrainStats();
 
+        public bool EmitTrainingAlignments { get; }
+
         public int MaxCorpusCount { get; set; } = int.MaxValue;
 
         public Task TrainAsync(IProgress<ProgressStatus> progress = null, CancellationToken cancellationToken = default)
@@ -243,6 +249,14 @@ void Report() =>
             Report();
             cancellationToken.ThrowIfCancellationRequested();
 
+            if (EmitTrainingAlignments)
+            {
+                // Retain the alignments computed during training so that they can be returned without a
+                // separate inference pass. Only the final (most refined) model's alignments are needed,
+                // since that is the model used for inference.
+                Thot.swAlignModel_setEmitTrainingAlignments(Handle, true);
+            }
+
             int trainedSegmentCount = 0;
             foreach ((IntPtr handle, int storedIterationCount) in _models)
             {

src/SIL.Machine/Corpora/CorporaExtensions.cs

@@ -1282,22 +1282,20 @@ public override IEnumerable<ParallelTextRow> GetRows(IEnumerable<string> textIds
             }
         }
 
-        private class WordAlignParallelTextCorpus : ParallelTextCorpusBase
+        private class WordAlignParallelTextCorpus : WordAlignParallelTextCorpusBase
         {
             private readonly IParallelTextCorpus _corpus;
             private readonly IWordAligner _aligner;
             private readonly int _batchSize;
 
             public WordAlignParallelTextCorpus(IParallelTextCorpus corpus, IWordAligner aligner, int batchSize)
+                : base(corpus)
             {
                 _corpus = corpus;
                 _aligner = aligner;
                 _batchSize = batchSize;
             }
 
-            public override bool IsSourceTokenized => _corpus.IsSourceTokenized;
-            public override bool IsTargetTokenized => _corpus.IsTargetTokenized;
-
             public override IEnumerable<ParallelTextRow> GetRows(IEnumerable<string> textIds)
             {
                 foreach (IReadOnlyList<ParallelTextRow> batch in _corpus.GetRows(textIds).Batch(_batchSize))

src/SIL.Machine/Corpora/WordAlignParallelTextCorpusBase.cs

@@ -0,0 +1,22 @@
+using System.Collections.Generic;
+
+namespace SIL.Machine.Corpora
+{
+    public abstract class WordAlignParallelTextCorpusBase : ParallelTextCorpusBase
+    {
+        private readonly IParallelTextCorpus _corpus;
+
+        protected WordAlignParallelTextCorpusBase(IParallelTextCorpus corpus)
+        {
+            _corpus = corpus;
+        }
+
+        public override bool IsSourceTokenized => _corpus.IsSourceTokenized;
+
+        public override bool IsTargetTokenized => _corpus.IsTargetTokenized;
+
+        public override int Count(bool includeEmpty = true, IEnumerable<string> textIds = null) =>
+            // Aligning does not add or remove rows, so counting need not align, which may train a model.
+            _corpus.Count(includeEmpty, textIds);
+    }
+}

src/SIL.Machine/Translation/ITransductiveWordAlignmentModel.cs

@@ -0,0 +1,8 @@
+namespace SIL.Machine.Translation
+{
+    public interface ITransductiveWordAlignmentModel
+    {
+        int TrainingAlignmentCount { get; }
+        WordAlignmentMatrix GetTrainingAlignment(int n);
+    }
+}