update: get_wandb_artifact

medrag_multi_modal/retrieval/common.py

@@ -1,7 +1,5 @@
 from enum import Enum
 
-import wandb
-
 
 class SimilarityMetric(Enum):
     COSINE = "cosine"
@@ -14,18 +12,6 @@ def mean_pooling(token_embeddings, mask):
     return sentence_embeddings
 
 
-def get_wandb_artifact(artifact_address: str, artifact_type: str):
-    if wandb.run:
-        artifact = wandb.run.use_artifact(artifact_address, type=artifact_type)
-        artifact_dir = artifact.download()
-    else:
-        api = wandb.Api()
-        artifact = api.artifact(artifact_address)
-        artifact_dir = artifact.download()
-    metadata = artifact.metadata
-    return artifact_dir, metadata
-
-
 def argsort_scores(scores: list[float], descending: bool = False):
     return [
         {"item": item, "original_index": idx}

medrag_multi_modal/retrieval/contriever_retrieval.py

@@ -15,7 +15,8 @@ from transformers import (
 
 import wandb
 
-from .common import SimilarityMetric, argsort_scores, get_wandb_artifact, mean_pooling
+from ..utils import get_wandb_artifact
+from .common import SimilarityMetric, argsort_scores, mean_pooling
 
 
 class ContrieverRetriever(weave.Model):
@@ -143,7 +144,7 @@ class ContrieverRetriever(weave.Model):
             and chunk dataset.
         """
         artifact_dir, metadata = get_wandb_artifact(
-            index_artifact_address, "contriever-index"
+            index_artifact_address, "contriever-index", get_metadata=True
         )
         with safetensors.torch.safe_open(
             os.path.join(artifact_dir, "vector_index.safetensors"), framework="pt"

medrag_multi_modal/utils.py

@@ -1,7 +1,9 @@
 import wandb
 
 
-def get_wandb_artifact(artifact_name: str, artifact_type: str) -> str:
+def get_wandb_artifact(
+    artifact_name: str, artifact_type: str, get_metadata: bool = False
+) -> str:
     if wandb.run:
         artifact = wandb.use_artifact(artifact_name, type=artifact_type)
         artifact_dir = artifact.download()
@@ -9,4 +11,6 @@ def get_wandb_artifact(artifact_name: str, artifact_type: str) -> str:
         api = wandb.Api()
         artifact = api.artifact(artifact_name)
         artifact_dir = artifact.download()
+    if get_metadata:
+        return artifact_dir, artifact.metadata
     return artifact_dir