update loading vlite with mismatch number of metadata

sdan · Apr 11, 2024 · 6d2bf9f · 6d2bf9f
1 parent 87d2687
commit 6d2bf9f
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 9 deletions.
diff --git a/tests/unit.py b/tests/unit.py
@@ -133,12 +133,6 @@ def tearDownClass(cls):
         print("\nTest times:")
         for test_name, test_time in cls.test_times.items():
             print(f"{test_name}: {test_time:.4f} seconds")
-        if os.path.exists('vlite-unit.npz'):
-            print("[+] Removing vlite")
-            os.remove('vlite-unit.npz')
-        if os.path.exists('vlite-unit.ctx'):
-            print("[+] Removing vlite")
-            os.remove('vlite-unit.ctx')
         if os.path.exists('contexts/vlite-unit.ctx'):
             print("[+] Removing vlite ctx")
             os.remove('contexts/vlite-unit.ctx')

diff --git a/vlite/main.py b/vlite/main.py
@@ -25,9 +25,9 @@ def __init__(self, collection=None, device='cpu', model_name='mixedbread-ai/mxba
             print("Number of metadata: ", len(ctx_file.metadata))
             self.index = {
                 chunk_id: {
-                    'text': ctx_file.contexts[idx],
+                    'text': ctx_file.contexts[idx] if idx < len(ctx_file.contexts) else "",
                     'metadata': ctx_file.metadata.get(chunk_id, {}),
-                    'binary_vector': np.array(ctx_file.embeddings[idx])
+                    'binary_vector': np.array(ctx_file.embeddings[idx]) if idx < len(ctx_file.embeddings) else np.zeros(self.model.embedding_size)
                 }
                 for idx, chunk_id in enumerate(ctx_file.metadata.keys())
             }
@@ -230,7 +230,8 @@ def save(self):
             for chunk_id, chunk_data in self.index.items():
                 ctx_file.add_embedding(chunk_data['binary_vector'])
                 ctx_file.add_context(chunk_data['text'])
-                ctx_file.add_metadata(chunk_id, chunk_data['metadata'])
+                if 'metadata' in chunk_data:
+                    ctx_file.add_metadata(chunk_id, chunk_data['metadata'])
         print("Collection saved successfully.")
 
     def clear(self):