Store subword indices for quality scores in Response

src/tests/common-impl.cpp

@@ -155,10 +155,11 @@ void TestSuite<Service>::qualityEstimatorWords(Ptr<TranslationModel> model) {
   std::string source = readFromStdin();
   const Response response = bridge_.translate(service_, model, std::move(source), responseOptions);
 
-  for (const auto &sentenceQualityEstimate : response.qualityScores) {
+  for (size_t sentenceIdx = 0; sentenceIdx < response.qualityScores.size(); ++sentenceIdx) {
+    const auto &sentenceQualityEstimate = response.qualityScores[sentenceIdx];
     std::cout << "[SentenceBegin]\n";
 
-    for (const auto &wordByteRange : sentenceQualityEstimate.wordByteRanges) {
+    for (const auto &wordByteRange : getWordByteRanges(response, sentenceIdx)) {
       const string_view word(response.target.text.data() + wordByteRange.begin, wordByteRange.size());
       std::cout << word << "\n";
     }

src/translator/definitions.h

@@ -42,6 +42,16 @@ struct ByteRange {
   bool operator==(ByteRange other) const { return begin == other.begin && end == other.end; }
 };
 
+/// A Subword range is mechanically the same as a `ByteRange`, but instead of
+/// describing a span of bytes, it describes a span of Subword tokens. Using
+/// `Annotation.word()` you can switch between the two.
+struct SubwordRange {
+  size_t begin;
+  size_t end;
+  const size_t size() const { return end - begin; }
+  bool operator==(SubwordRange other) const { return begin == other.begin && end == other.end; }
+};
+
 class Response;
 using CallbackType = std::function<void(Response &&)>;
 

src/translator/quality_estimator.cpp

@@ -27,7 +27,7 @@ Response::SentenceQualityScore UnsupervisedQualityEstimator::computeSentenceScor
   const float sentenceScore =
       std::accumulate(std::begin(wordScores), std::end(wordScores), float(0.0)) / wordScores.size();
 
-  return {wordScores, subwordToWords(wordIndices, target, sentenceIdx), sentenceScore};
+  return {wordScores, wordIndices, sentenceScore};
 }
 
 LogisticRegressorQualityEstimator::Matrix::Matrix(const size_t rowsParam, const size_t colsParam)
@@ -160,7 +160,7 @@ Response::SentenceQualityScore LogisticRegressorQualityEstimator::computeSentenc
   const float sentenceScore =
       std::accumulate(std::begin(wordScores), std::end(wordScores), float(0.0)) / wordScores.size();
 
-  return {wordScores, subwordToWords(wordIndices, target, sentenceIdx), sentenceScore};
+  return {wordScores, wordIndices, sentenceScore};
 }
 
 std::vector<float> LogisticRegressorQualityEstimator::predict(const Matrix& features) const {
@@ -267,22 +267,4 @@ std::vector<SubwordRange> mapWords(const std::vector<float>& logProbs, const Ann
   return wordIndices;
 }
 
-std::vector<ByteRange> subwordToWords(const std::vector<SubwordRange>& wordIndices, const AnnotatedText& target,
-                                      const size_t sentenceIdx) {
-  std::vector<ByteRange> words;
-
-  for (const SubwordRange& wordIndice : wordIndices) {
-    size_t wordBegin = target.wordAsByteRange(sentenceIdx, wordIndice.begin).begin;
-    size_t wordEnd = target.wordAsByteRange(sentenceIdx, wordIndice.end).begin;
-
-    if (isspace(target.text.at(wordBegin))) {
-      ++wordBegin;
-    }
-
-    words.emplace_back(ByteRange{wordBegin, wordEnd});
-  }
-
-  return words;
-}
-
 }  // namespace marian::bergamot

src/translator/quality_estimator.h

@@ -21,8 +21,6 @@ class QualityEstimator {
   virtual void computeQualityScores(const Histories &histories, Response &response) const = 0;
 };
 
-using SubwordRange = ByteRange;
-
 /// Unsupervised Quality Estimator model. It uses the translator model's log probabilities (log probs) as a proxy for
 /// quality scores. Then, for a given word, its quality score is computed by taking the mean of the log probs of the
 /// tokens that make it up. The sentence score is the mean of all word's log probs.
@@ -209,14 +207,4 @@ inline std::shared_ptr<QualityEstimator> createQualityEstimator(const AlignedMem
 std::vector<SubwordRange> mapWords(const std::vector<float> &logProbs, const AnnotatedText &target,
                                    const size_t sentenceIdx);
 
-/// Given a vector of subwordRanges, it maps the elements to be real words rather than sublevel tokens. The words are
-/// represented through ByteRanges.
-
-/// @param [in] wordIndices: A vector where each element correspond to the index of a real word and its values are
-/// represented by the SubwordRanges (which are aliases of ByteRanges) which represents sublevel token positions
-/// @param [in] target: AnnotatedText target value
-/// @param [in] sentenceIdx: the id of a candidate sentence
-std::vector<ByteRange> subwordToWords(const std::vector<SubwordRange> &wordIndices, const AnnotatedText &target,
-                                      const size_t sentenceIdx);
-
 }  // namespace marian::bergamot

src/translator/response.cpp

@@ -142,4 +142,22 @@ std::vector<Alignment> remapAlignments(const Response &first, const Response &se
   return alignments;
 }
 
+std::vector<ByteRange> getWordByteRanges(const Response &response, size_t sentenceIdx) {
+  std::vector<ByteRange> wordByteRanges;
+  wordByteRanges.reserve(response.qualityScores[sentenceIdx].wordRanges.size());
+
+  for (auto &&word : response.qualityScores[sentenceIdx].wordRanges) {
+    size_t wordBegin = response.target.wordAsByteRange(sentenceIdx, word.begin).begin;
+    size_t wordEnd = response.target.wordAsByteRange(sentenceIdx, word.end).begin;
+
+    if (std::isspace(response.target.text.at(wordBegin))) {
+      ++wordBegin;
+    }
+
+    wordByteRanges.emplace_back(ByteRange{wordBegin, wordEnd});
+  }
+
+  return wordByteRanges;
+}
+
 }  // namespace marian::bergamot

src/translator/response.h

@@ -30,8 +30,8 @@ struct Response {
   struct SentenceQualityScore {
     /// Quality score of each translated word
     std::vector<float> wordScores;
-    /// Each word position in the translated text
-    std::vector<ByteRange> wordByteRanges;
+    /// Position of start and end token of each word in the translated text
+    std::vector<SubwordRange> wordRanges;
     /// Whole sentence quality score (it is composed by the mean of its words)
     float sentenceScore = 0.0;
   };
@@ -77,6 +77,8 @@ struct Response {
 
 std::vector<Alignment> remapAlignments(const Response &first, const Response &second);
 
+std::vector<ByteRange> getWordByteRanges(Response const &response, size_t sentenceIdx);
+
 }  // namespace bergamot
 }  // namespace marian
 

wasm/bindings/response_bindings.cpp

@@ -10,17 +10,38 @@
 #include "response.h"
 
 using Response = marian::bergamot::Response;
-using SentenceQualityScore = marian::bergamot::Response::SentenceQualityScore;
 using ByteRange = marian::bergamot::ByteRange;
 
+/// Same type as Response::SentenceQualityScore, except with wordByteRanges
+/// instead of wordRanges.
+struct SentenceQualityScore {
+  /// Quality score of each translated word
+  std::vector<float> wordScores;
+  /// Position of each word in the translated text
+  std::vector<ByteRange> wordByteRanges;
+  /// Whole sentence quality score (it is composed by the mean of its words)
+  float sentenceScore = 0.0;
+};
+
 using namespace emscripten;
 
 // Binding code
 EMSCRIPTEN_BINDINGS(byte_range) {
   value_object<ByteRange>("ByteRange").field("begin", &ByteRange::begin).field("end", &ByteRange::end);
 }
 
-std::vector<SentenceQualityScore> getQualityScores(const Response& response) { return response.qualityScores; }
+std::vector<SentenceQualityScore> getQualityScores(const Response& response) {
+  std::vector<SentenceQualityScore> scores;
+  scores.reserve(response.qualityScores.size());
+
+  for (size_t sentenceIdx = 0; sentenceIdx < response.qualityScores.size(); ++sentenceIdx) {
+    scores.emplace_back(SentenceQualityScore{response.qualityScores[sentenceIdx].wordScores,
+                                             marian::bergamot::getWordByteRanges(response, sentenceIdx),
+                                             response.qualityScores[sentenceIdx].sentenceScore});
+  }
+
+  return scores;
+}
 
 EMSCRIPTEN_BINDINGS(response) {
   class_<Response>("Response")

wasm/test_page/js/worker.js

@@ -326,6 +326,7 @@ const _parseTranslatedTextSentenceQualityScores = (vectorResponse) => {
       sentenceQualityScores.push(sentenceQualityScore);
     }
     result.push(sentenceQualityScores);
+    vectorSentenceQualityScore.delete();
   }
   return result;
 }