Fixes sillsdev/serval#202

sillsdev · Nov 7, 2023 · fba7aea · fba7aea
1 parent de12933
commit fba7aea
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/src/SIL.Machine.AspNetCore/Models/Corpus.cs b/src/SIL.Machine.AspNetCore/Models/Corpus.cs
@@ -5,7 +5,9 @@ public class Corpus
     public string Id { get; set; } = default!;
     public string SourceLanguage { get; set; } = default!;
     public string TargetLanguage { get; set; } = default!;
+    public bool TrainOnAll { get; set; }
     public bool PretranslateAll { get; set; }
+    public HashSet<string> TrainOnTextIds { get; set; } = default!;
     public HashSet<string> PretranslateTextIds { get; set; } = default!;
     public List<CorpusFile> SourceFiles { get; set; } = default!;
     public List<CorpusFile> TargetFiles { get; set; } = default!;

diff --git a/src/SIL.Machine.AspNetCore/Services/NmtPreprocessBuildJob.cs b/src/SIL.Machine.AspNetCore/Services/NmtPreprocessBuildJob.cs
@@ -76,8 +76,11 @@ async IAsyncEnumerable<Pretranslation> ProcessRowsAsync()
 
                 foreach (ParallelTextRow row in parallelCorpus)
                 {
-                    await sourceTrainWriter.WriteAsync($"{row.SourceText}\n");
-                    await targetTrainWriter.WriteAsync($"{row.TargetText}\n");
+                    if (corpus.TrainOnAll || corpus.TrainOnTextIds.Contains(row.TextId))
+                    {
+                        await sourceTrainWriter.WriteAsync($"{row.SourceText}\n");
+                        await targetTrainWriter.WriteAsync($"{row.TargetText}\n");
+                    }
                     if (
                         (corpus.PretranslateAll || corpus.PretranslateTextIds.Contains(row.TextId))
                         && row.SourceSegment.Count > 0

diff --git a/src/SIL.Machine.AspNetCore/Services/ServalTranslationEngineServiceV1.cs b/src/SIL.Machine.AspNetCore/Services/ServalTranslationEngineServiceV1.cs
@@ -236,7 +236,9 @@ private static Models.Corpus Map(Serval.Translation.V1.Corpus source)
             Id = source.Id,
             SourceLanguage = source.SourceLanguage,
             TargetLanguage = source.TargetLanguage,
+            TrainOnAll = source.TrainOnAll,
             PretranslateAll = source.PretranslateAll,
+            TrainOnTextIds = source.TrainOnTextIds.ToHashSet(),
             PretranslateTextIds = source.PretranslateTextIds.ToHashSet(),
             SourceFiles = source.SourceFiles.Select(Map).ToList(),
             TargetFiles = source.TargetFiles.Select(Map).ToList()