Merge pull request #4763 from vespa-engine/bratseth/support-binary-tensor-files

Support binary tensor files with ending .tbf

3 files changed, 13 insertions, 2 deletions

diff --git a/eval/src/tests/eval/value_cache/dense.tbf b/eval/src/tests/eval/value_cache/dense.tbf
new file mode 100644
index 00000000000..61012a35ff9
--- /dev/null
+++ b/eval/src/tests/eval/value_cache/dense.tbf
diff --git a/eval/src/tests/eval/value_cache/tensor_loader_test.cpp b/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
index 5dd8caa6e27..8180a7daef8 100644
--- a/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
+++ b/eval/src/tests/eval/value_cache/tensor_loader_test.cpp
@@ -67,6 +67,10 @@ TEST_F("require that lz4 compressed dense tensor can be loaded", ConstantTensorL
     TEST_DO(verify_tensor(make_dense_tensor(), f1.create(TEST_PATH("dense.json.lz4"), "tensor(x[2],y[2])")));
 }
 
+TEST_F("require that a binary tensor can be loaded", ConstantTensorLoader(SimpleTensorEngine::ref())) {
+    TEST_DO(verify_tensor(make_dense_tensor(), f1.create(TEST_PATH("dense.tbf"), "tensor(x[2],y[2])")));
+}
+
 TEST_F("require that lz4 compressed sparse tensor can be loaded", ConstantTensorLoader(SimpleTensorEngine::ref())) {
     TEST_DO(verify_tensor(make_sparse_tensor(), f1.create(TEST_PATH("sparse.json.lz4"), "tensor(x{},y{})")));
 }
diff --git a/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp b/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp
index 38d5bbc643b..afc8471bdb4 100644
--- a/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp
+++ b/eval/src/vespa/eval/eval/value_cache/constant_tensor_loader.cpp
@@ -4,6 +4,7 @@
 #include <vespa/eval/eval/tensor.h>
 #include <vespa/eval/eval/tensor_engine.h>
 #include <vespa/eval/eval/tensor_spec.h>
+#include <vespa/vespalib/objects/nbostream.h>
 #include <vespa/vespalib/io/mapped_file_input.h>
 #include <vespa/vespalib/data/lz4_input_decoder.h>
 #include <vespa/vespalib/data/slime/slime.h>
@@ -54,14 +55,14 @@ void decode_json(const vespalib::string &path, Slime &slime) {
     } else {
         if (ends_with(path, ".lz4")) {
             size_t buffer_size = 64 * 1024;
-            Lz4InputDecoder lz4_decoder(file, buffer_size);            
+            Lz4InputDecoder lz4_decoder(file, buffer_size);
             decode_json(path, lz4_decoder, slime);
             if (lz4_decoder.failed()) {
                 LOG(warning, "file contains lz4 errors (%s): %s",
                     lz4_decoder.reason().c_str(), path.c_str());
             }
         } else {
-            decode_json(path, file, slime);            
+            decode_json(path, file, slime);
         }
     }
 }
@@ -76,6 +77,12 @@ ConstantTensorLoader::create(const vespalib::string &path, const vespalib::strin
         LOG(warning, "invalid type specification: %s", type.c_str());
         return std::make_unique<SimpleConstantValue>(_engine.from_spec(TensorSpec("double")));
     }
+    if (ends_with(path, ".tbf")) {
+        vespalib::MappedFileInput file(path);
+        vespalib::Memory content = file.get();
+        vespalib::nbostream stream(content.data, content.size);
+        return std::make_unique<SimpleConstantValue>(_engine.decode(stream));
+    }
     Slime slime;
     decode_json(path, slime);
     std::set<vespalib::string> indexed;