From d2c7cfa8e50cfd5d8dbee9fe4e45f82474c0b6fd Mon Sep 17 00:00:00 2001
From: Alexander Gutkin <agutkin@google.com>
Date: Wed, 11 Sep 2024 08:09:41 -0700
Subject: [PATCH] Parallelizing the decoder FST construction.

PiperOrigin-RevId: 673397357
---
 nisaba/translit/fst/pairlm_decoder.cc | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/nisaba/translit/fst/pairlm_decoder.cc b/nisaba/translit/fst/pairlm_decoder.cc
index 7451cb3c..49d46159 100644
--- a/nisaba/translit/fst/pairlm_decoder.cc
+++ b/nisaba/translit/fst/pairlm_decoder.cc
@@ -993,11 +993,18 @@ StdVectorFst PairLMDecoder::BuildTransliterationFst(
   const std::vector<std::string> input_words = absl::StrSplit(
       input_line, utf8::Utf8WhitespaceDelimiter(), absl::SkipEmpty());
   std::vector<StdVectorFst> translit_fsts(input_words.size());
-     // TODO: Fix parallelization.
+  {
+     ThreadPool pool(max_parallel_tokens_);
+     pool.StartWorkers();
      for (int i = 0; i < input_words.size(); ++i) {
-       translit_fsts[i] = TransliterateWord(input_words[i], word_k_best,
-                                            fst_params);
+       pool.Schedule([this, i, word_k_best, &translit_fsts, &input_words,
+         &fst_params] {
+         translit_fsts[i] = TransliterateWord(input_words[i], word_k_best,
+                                              fst_params);
+       });
      }
+     // Wait until all threads are complete.
+  }
 
   // Loop through the fsts produced, joining them into a sentence lattice.
   for (int i = 0; i < input_words.size(); ++i) {