Port transductive alignment model

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

@@ -0,0 +1,108 @@
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using SIL.Machine.Corpora;
+using SIL.Machine.Utils;
+
+namespace SIL.Machine.Translation.Thot
+{
+    public static class IParallelTextCorpusExtensions
+    {
+        public static async Task<IParallelTextCorpus> WordAlignAsync(
+            this IParallelTextCorpus corpus,
+            ThotWordAlignmentModelType modelType,
+            int batchSize = 1024,
+            SymmetrizationHeuristic symmetrizationHeuristic = SymmetrizationHeuristic.GrowDiagFinalAnd,
+            IProgress<ProgressStatus> progress = null
+        )
+        {
+            // TODO: Warning! This model must be disposed correctly.
+            var model = new ThotSymmetrizedWordAlignmentModel(
+                ThotWordAlignmentModel.Create(modelType),
+                ThotWordAlignmentModel.Create(modelType)
+            )
+            {
+                Heuristic = symmetrizationHeuristic,
+                // Retain the alignments computed during training so that the corpus can be aligned
+                // without a separate, potentially expensive, inference pass.
+                EmitTrainingAlignments = true,
+            };
+
+            using (ITrainer trainer = model.CreateTrainer(corpus))
+            {
+                await trainer.TrainAsync(progress);
+                await trainer.SaveAsync();
+            }
+
+            return corpus.WordAlign(model, batchSize);
+        }
+
+        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 : ParallelTextCorpusBase
+        {
+            private readonly IParallelTextCorpus _corpus;
+            private readonly ITransductiveWordAlignmentModel _model;
+
+            public TransductiveWordAlignParallelTextCorpus(
+                IParallelTextCorpus corpus,
+                ITransductiveWordAlignmentModel model
+            )
+            {
+                _corpus = corpus;
+                _model = model;
+            }
+
+            public override bool IsSourceTokenized => _corpus.IsSourceTokenized;
+
+            public override bool IsTargetTokenized => _corpus.IsTargetTokenized;
+
+            public override IEnumerable<ParallelTextRow> GetRows(IEnumerable<string> textIds)
+            {
+                // The training alignments are keyed by the order in which the sentence pairs were added
+                // during training, so the full corpus must be iterated to keep the index in sync; rows that
+                // are not in the requested texts are skipped rather than filtered out of the enumeration.
+                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,56 @@
+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 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/Translation/ITransductiveWordAlignmentModel.cs

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