Support tie_embeddings from Nematus: CPU

1 files changed, 23 insertions, 0 deletions

diff --git a/src/amun/cpu/npz_converter.h b/src/amun/cpu/npz_converter.h
index 7488ead7..ffb7c692 100644
--- a/src/amun/cpu/npz_converter.h
+++ b/src/amun/cpu/npz_converter.h
@@ -79,6 +79,29 @@ class NpzConverter {
       return std::move(ret);
     }
 
+    mblas::Matrix operator[](const std::vector<std::pair<std::string, bool>> keys) const {
+      BlazeWrapper matrix;
+      for (auto key : keys) {
+        auto it = model_.find(key.first);
+        if(it != model_.end()) {
+          NpyMatrixWrapper np(it->second);
+          matrix = BlazeWrapper(np.data(), np.size1(), np.size2());
+          mblas::Matrix ret;
+          if (key.second) {
+            const auto matrix2 = blaze::trans(matrix);
+            ret = matrix2;
+          } else {
+            ret = matrix;
+          }
+          return std::move(ret);
+        }
+      }
+      std::cerr << "Matrix not found: " << keys[0].first << "\n";
+
+      mblas::Matrix ret;
+      return std::move(ret);
+    }
+
     mblas::Matrix operator()(const std::string& key,
                                    bool transpose) const {
       BlazeWrapper matrix;