geekyrakshit commited on
Commit
05b69a5
·
1 Parent(s): af715dd

fix: bug in BM25sRetriever

Browse files
medrag_multi_modal/retrieval/bm25s_retrieval.py CHANGED
@@ -104,6 +104,20 @@ class BM25sRetriever(weave.Model):
104
  initialize the class instance with the appropriate language and stemming
105
  settings.
106
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
107
  Args:
108
  index_artifact_address (str): The address of the Weights & Biases artifact
109
  containing the BM25 index.
@@ -121,8 +135,9 @@ class BM25sRetriever(weave.Model):
121
  api = wandb.Api()
122
  artifact = api.artifact(index_artifact_address)
123
  artifact_dir = artifact.download()
124
- index_name = glob(os.path.join(artifact_dir, "*"))[0].split("/")[-1]
125
- retriever = bm25s.BM25.load(index_name, load_corpus=True)
 
126
  metadata = artifact.metadata
127
  return cls(
128
  language=metadata["language"],
@@ -183,7 +198,7 @@ class BM25sRetriever(weave.Model):
183
  load_dotenv()
184
  weave.init(project_name="ml-colabs/medrag-multi-modal")
185
  retriever = BM25sRetriever.from_wandb_artifact(
186
- index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-bm25s:v2"
187
  )
188
  retrieved_chunks = retriever.predict(query="What are Ribosomes?")
189
  ```
 
104
  initialize the class instance with the appropriate language and stemming
105
  settings.
106
 
107
+ !!! example "Example Usage"
108
+ ```python
109
+ import weave
110
+ from dotenv import load_dotenv
111
+
112
+ from medrag_multi_modal.retrieval import BM25sRetriever
113
+
114
+ load_dotenv()
115
+ weave.init(project_name="ml-colabs/medrag-multi-modal")
116
+ retriever = BM25sRetriever.from_wandb_artifact(
117
+ index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-bm25s:latest"
118
+ )
119
+ ```
120
+
121
  Args:
122
  index_artifact_address (str): The address of the Weights & Biases artifact
123
  containing the BM25 index.
 
135
  api = wandb.Api()
136
  artifact = api.artifact(index_artifact_address)
137
  artifact_dir = artifact.download()
138
+ retriever = bm25s.BM25.load(
139
+ glob(os.path.join(artifact_dir, "*"))[0], load_corpus=True
140
+ )
141
  metadata = artifact.metadata
142
  return cls(
143
  language=metadata["language"],
 
198
  load_dotenv()
199
  weave.init(project_name="ml-colabs/medrag-multi-modal")
200
  retriever = BM25sRetriever.from_wandb_artifact(
201
+ index_artifact_address="ml-colabs/medrag-multi-modal/grays-anatomy-bm25s:latest"
202
  )
203
  retrieved_chunks = retriever.predict(query="What are Ribosomes?")
204
  ```