add: utility save_vector_index

medrag_multi_modal/retrieval/common.py

@@ -1,5 +1,11 @@
 from enum import Enum
 
+import safetensors
+import safetensors.torch
+import torch
+
+import wandb
+
 
 class SimilarityMetric(Enum):
     COSINE = "cosine"
@@ -19,3 +25,20 @@ def argsort_scores(scores: list[float], descending: bool = False):
             list(enumerate(scores)), key=lambda x: x[1], reverse=descending
         )
     ]
+
+
+def save_vector_index(
+    vector_index: torch.Tensor,
+    index_name: str,
+    metadata: dict,
+    filename: str = "vector_index.safetensors",
+):
+    safetensors.torch.save_file({"vector_index": vector_index.cpu()}, filename)
+    if wandb.run:
+        artifact = wandb.Artifact(
+            name=index_name,
+            type="contriever-index",
+            metadata=metadata,
+        )
+        artifact.add_file(filename)
+        artifact.save()