Export alignments through WASM

For browsermt#362.
jerinphilip · Mar 21, 2022 · 60e88b6 · 60e88b6
1 parent 409b7d2
commit 60e88b6
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 1 deletion.
diff --git a/wasm/bindings/response_bindings.cpp b/wasm/bindings/response_bindings.cpp
@@ -12,6 +12,10 @@
 using Response = marian::bergamot::Response;
 using ByteRange = marian::bergamot::ByteRange;
 
+// Discrete Probability Distribution, named Distribution for brevity.
+using Distribution = std::vector<float>;
+using Alignment = std::vector<Distribution>;
+
 using namespace emscripten;
 
 // Binding code
@@ -20,13 +24,18 @@ EMSCRIPTEN_BINDINGS(byte_range) {
 }
 
 EMSCRIPTEN_BINDINGS(response) {
+  register_vector<float>("Distribution");
+  register_vector<Distribution>("Alignment");
+  register_vector<Alignment>("Alignments");
+
   class_<Response>("Response")
       .constructor<>()
       .function("size", &Response::size)
       .function("getOriginalText", &Response::getOriginalText)
       .function("getTranslatedText", &Response::getTranslatedText)
       .function("getSourceSentence", &Response::getSourceSentenceAsByteRange)
-      .function("getTranslatedSentence", &Response::getTargetSentenceAsByteRange);
+      .function("getTranslatedSentence", &Response::getTargetSentenceAsByteRange)
+      .property("alignments", &Response::alignments);
 
   register_vector<Response>("VectorResponse");
 }
diff --git a/wasm/test_page/js/worker.js b/wasm/test_page/js/worker.js
@@ -150,6 +150,7 @@ const translate = (from, to, input, translateOptions) => {
     const listSourceText = _parseSourceText(vectorResponse);
     const listTranslatedTextSentences = _parseTranslatedTextSentences(vectorResponse);
     const listSourceTextSentences = _parseSourceTextSentences(vectorResponse);
+    const listAlignments = _parseAlignments(vectorResponse);
 
     log(`Source text: ${listSourceText}`);
     log(`Translated text: ${listTranslatedText}`);
@@ -253,6 +254,18 @@ const _getLoadedTranslationModel = (srcLang, tgtLang) => {
   return languagePairToTranslationModels.get(languagePair);
 }
 
+const _parseAlignments = (vectorResponse) => {
+  const result = [];
+  for (let i = 0; i < vectorResponse.size(); i++) {
+    const response = vectorResponse.get(i);
+    log("Parsing Alignments");
+    result.push(response.alignments);
+    log(response.alignments);
+  }
+  return result;
+}
+
+
 const _parseTranslatedText = (vectorResponse) => {
   const result = [];
   for (let i = 0; i < vectorResponse.size(); i++) {