Bundle AlignedMemory inputs with MemoryBundle (#147)

app/service-cli.cpp

@@ -16,19 +16,15 @@ int main(int argc, char *argv[]) {
   auto cp = marian::bergamot::createConfigParser();
   auto options = cp.parseOptions(argc, argv, true);
 
-  // Prepare memories for model and shortlist
-  marian::bergamot::AlignedMemory modelBytes, shortlistBytes;
-  std::vector<std::shared_ptr<marian::bergamot::AlignedMemory>> vocabsBytes;
+  // Prepare memories for bytearrays (including model, shortlist and vocabs)
+  marian::bergamot::MemoryBundle memoryBundle;
 
   if (options->get<bool>("check-bytearray")) {
     // Load legit values into bytearrays.
-    modelBytes = marian::bergamot::getModelMemoryFromConfig(options);
-    shortlistBytes = marian::bergamot::getShortlistMemoryFromConfig(options);
-    marian::bergamot::getVocabsMemoryFromConfig(options, vocabsBytes);
+    memoryBundle = marian::bergamot::getMemoryBundleFromConfig(options);
   }
 
-  marian::bergamot::Service service(options, std::move(modelBytes),
-                                    std::move(shortlistBytes), std::move(vocabsBytes));
+  marian::bergamot::Service service(options, std::move(memoryBundle));
 
   // Read a large input text blob from stdin
   std::ostringstream std_input;

src/translator/byte_array_util.cpp

@@ -117,5 +117,13 @@ void getVocabsMemoryFromConfig(marian::Ptr<marian::Options> options,
   }
 }
 
+MemoryBundle getMemoryBundleFromConfig(marian::Ptr<marian::Options> options){
+  MemoryBundle memoryBundle;
+  memoryBundle.model = getModelMemoryFromConfig(options);
+  memoryBundle.shortlist = getShortlistMemoryFromConfig(options);
+  getVocabsMemoryFromConfig(options, memoryBundle.vocabs);
+  return memoryBundle;
+}
+
 } // namespace bergamot
 } // namespace marian

src/translator/byte_array_util.h

@@ -10,5 +10,6 @@ AlignedMemory getShortlistMemoryFromConfig(marian::Ptr<marian::Options> options)
 void getVocabsMemoryFromConfig(marian::Ptr<marian::Options> options,
                                std::vector<std::shared_ptr<AlignedMemory>>& vocabMemories);
 bool validateBinaryModel(const AlignedMemory& model, uint64_t fileSize);
+MemoryBundle getMemoryBundleFromConfig(marian::Ptr<marian::Options> options);
 } // namespace bergamot
 } // namespace marian

src/translator/definitions.h

@@ -15,6 +15,42 @@ typedef std::vector<Segment> Segments;
 /// Shortcut to AlignedVector<char> for byte arrays
 typedef AlignedVector<char> AlignedMemory;
 
+/// Memory bundle for all byte-arrays.
+/// Can be a set/subset of model, shortlist, vocabs and ssplitPrefixFile bytes.
+struct MemoryBundle {
+  AlignedMemory model;  ///< Byte-array of model (aligned to 256)
+  AlignedMemory shortlist;  ///< Byte-array of shortlist (aligned to 64)
+
+  /// Vector of vocabulary memories (aligned to 64).
+  /// If two vocabularies are the same (based on the filenames), two entries (shared
+  /// pointers) will be generated which share the same AlignedMemory object.
+  std::vector<std::shared_ptr<AlignedMemory>> vocabs;
+
+  /// @todo Not implemented yet
+  AlignedMemory ssplitPrefixFile;
+
+  MemoryBundle() = default;
+
+  MemoryBundle(MemoryBundle &&from){
+    model = std::move(from.model);
+    shortlist = std::move(from.shortlist);
+    vocabs = std::move(vocabs);
+    ssplitPrefixFile = std::move(from.ssplitPrefixFile);
+  }
+
+  MemoryBundle &operator=(MemoryBundle &&from) {
+    model = std::move(from.model);
+    shortlist = std::move(from.shortlist);
+    vocabs = std::move(vocabs);
+    ssplitPrefixFile = std::move(from.ssplitPrefixFile);
+    return *this;
+  }
+
+  // Delete copy constructors
+  MemoryBundle(const MemoryBundle&) = delete;
+  MemoryBundle& operator=(const MemoryBundle&) = delete;
+};
+
 } // namespace bergamot
 } // namespace marian
 

src/translator/service.cpp

@@ -41,14 +41,13 @@ loadVocabularies(marian::Ptr<marian::Options> options,
 namespace marian {
 namespace bergamot {
 
-Service::Service(Ptr<Options> options, AlignedMemory modelMemory, AlignedMemory shortlistMemory,
-                 std::vector<std::shared_ptr<AlignedMemory>> vocabMemories)
+Service::Service(Ptr<Options> options, MemoryBundle memoryBundle)
     : requestId_(0), options_(options),
-      vocabs_(std::move(loadVocabularies(options, std::move(vocabMemories)))),
+      vocabs_(std::move(loadVocabularies(options, std::move(memoryBundle.vocabs)))),
       text_processor_(vocabs_, options), batcher_(options),
       numWorkers_(options->get<int>("cpu-threads")),
-      modelMemory_(std::move(modelMemory)),
-      shortlistMemory_(std::move(shortlistMemory))
+      modelMemory_(std::move(memoryBundle.model)),
+      shortlistMemory_(std::move(memoryBundle.shortlist))
 #ifndef WASM_COMPATIBLE_SOURCE
       // 0 elements in PCQueue is illegal and can lead to failures. Adding a
       // guard to have at least one entry allocated. In the single-threaded

src/translator/service.h

@@ -55,53 +55,29 @@ namespace bergamot {
 /// // Do things with response.
 /// ```
 ///
-/// Optionally Service can be initialized by also passing model memory for
-/// purposes of efficiency (which defaults to nullpointer and then reads from
+/// Optionally Service can be initialized by also passing bytearray memories
+/// for purposes of efficiency (which defaults to empty and then reads from
 /// file supplied through config).
 ///
 class Service {
 
 public:
+  /// Construct Service from Marian options. If memoryBundle is empty, Service is
+  /// initialized from file-based loading. Otherwise, Service is initialized from
+  /// the given bytearray memories.
   /// @param options Marian options object
-  /// @param modelMemory byte array (aligned to 256!!!) that contains the bytes
-  /// of a model.bin.
-  /// @param shortlistMemory byte array of shortlist (aligned to 64)
-  /// @param vocabMemories vector of vocabulary memories (aligned to 64)
-  explicit Service(Ptr<Options> options, AlignedMemory modelMemory,
-                   AlignedMemory shortlistMemory,
-                   std::vector<std::shared_ptr<AlignedMemory>> vocabMemories);
+  /// @param memoryBundle holds all byte-array memories. Can be a set/subset of
+  /// model, shortlist, vocabs and ssplitPrefixFile bytes. Optional.
+  explicit Service(Ptr<Options> options, MemoryBundle memoryBundle={});
 
-  /// Construct Service purely from Options. This expects options which
-  /// marian-decoder expects to be set for loading model shortlist and
-  /// vocabularies from files in addition to parameters that set unset desired
-  /// features (e.g: alignments, quality-scores).
-  ///
-  /// This is equivalent to a call to:
-  /// ```cpp
-  ///    Service(options, AlignedMemory(), AlignedMemory(), {})
-  /// ```
-  /// wherein empty memory is passed and internal flow defaults to file-based
-  /// model, shortlist loading. AlignedMemory() corresponds to empty memory
-  explicit Service(Ptr<Options> options)
-      : Service(options, AlignedMemory(), AlignedMemory(), {}) {}
-
-  /// Construct Service from a string configuration.
-  /// @param [in] config string parsable as YAML expected to adhere with marian
-  /// config
-  /// @param [in] modelMemory byte array (aligned to 256!!!) that contains the
-  /// bytes of a model.bin. Optional. AlignedMemory() corresponds to empty memory
-  /// @param [in] shortlistMemory byte array of shortlist (aligned to 64). Optional.
-  /// @param [in] vocabMemories vector of vocabulary memories (aligned to 64). Optional.
-  /// If two vocabularies are the same (based on the filenames), two entries (shared
-  /// pointers) will be generated which share the same AlignedMemory object.
-  explicit Service(const std::string &config,
-                   AlignedMemory modelMemory = AlignedMemory(),
-                   AlignedMemory shortlistMemory = AlignedMemory(),
-                   std::vector<std::shared_ptr<AlignedMemory>> vocabsMemories = {})
-      : Service(parseOptions(config, /*validate=*/false),
-                std::move(modelMemory),
-                std::move(shortlistMemory),
-                std::move(vocabsMemories)) {}
+  /// Construct Service from a string configuration. If memoryBundle is empty, Service is
+  /// initialized from file-based loading. Otherwise, Service is initialized from
+  /// the given bytearray memories.
+  /// @param [in] config string parsable as YAML expected to adhere with marian config
+  /// @param [in] memoryBundle holds all byte-array memories. Can be a set/subset of
+  /// model, shortlist, vocabs and ssplitPrefixFile bytes. Optional.
+  explicit Service(const std::string &config, MemoryBundle memoryBundle={})
+      : Service(parseOptions(config, /*validate=*/false), std::move(memoryBundle)) {}
 
   /// Explicit destructor to clean up after any threads initialized in
   /// asynchronous operation mode.

wasm/bindings/TranslationModelBindings.cpp

@@ -48,14 +48,22 @@ std::vector<std::shared_ptr<AlignedMemory>> prepareVocabsSmartMemories(std::vect
   return vocabsSmartMemories;
 }
 
+marian::bergamot::MemoryBundle prepareMemoryBundle(AlignedMemory* modelMemory,
+                                                   AlignedMemory* shortlistMemory,
+                                                   std::vector<AlignedMemory*> uniqueVocabsMemories){
+  marian::bergamot::MemoryBundle memoryBundle;
+  memoryBundle.model = std::move(*modelMemory);
+  memoryBundle.shortlist = std::move(*shortlistMemory);
+  memoryBundle.vocabs = std::move(prepareVocabsSmartMemories(uniqueVocabsMemories));
+
+  return memoryBundle;
+}
+
 TranslationModel* TranslationModelFactory(const std::string &config,
                                           AlignedMemory* modelMemory,
                                           AlignedMemory* shortlistMemory,
                                           std::vector<AlignedMemory*> uniqueVocabsMemories) {
-  return new TranslationModel(config,
-                              std::move(*modelMemory),
-                              std::move(*shortlistMemory),
-                              std::move(prepareVocabsSmartMemories(uniqueVocabsMemories)));
+  return new TranslationModel(config, std::move(prepareMemoryBundle(modelMemory, shortlistMemory, uniqueVocabsMemories)));
 }
 
 EMSCRIPTEN_BINDINGS(translation_model) {