geekyrakshit commited on
Commit
88a5bcf
·
1 Parent(s): 4ea2b30

update: BM25sRetriever

Browse files
medrag_multi_modal/retrieval/bm25s_retrieval.py CHANGED
@@ -1,3 +1,5 @@
 
 
1
  from typing import Optional
2
 
3
  import bm25s
@@ -36,10 +38,51 @@ class BM25sRetriever(weave.Model):
36
  stemmer=Stemmer(self.language) if self.use_stemmer else None,
37
  )
38
  self._retriever.index(corpus_tokens)
39
- self._retriever.save(index_name, corpus=[dict(row) for row in corpus_dataset])
40
  if index_name:
41
- self._retriever.save(index_name)
 
 
42
  if wandb.run:
43
- artifact = wandb.Artifact(name=index_name, type="bm25s-index")
44
- artifact.add_dir(index_name)
 
 
 
 
 
 
 
45
  artifact.save()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from glob import glob
3
  from typing import Optional
4
 
5
  import bm25s
 
38
  stemmer=Stemmer(self.language) if self.use_stemmer else None,
39
  )
40
  self._retriever.index(corpus_tokens)
 
41
  if index_name:
42
+ self._retriever.save(
43
+ index_name, corpus=[dict(row) for row in corpus_dataset]
44
+ )
45
  if wandb.run:
46
+ artifact = wandb.Artifact(
47
+ name=index_name,
48
+ type="bm25s-index",
49
+ metadata={
50
+ "language": self.language,
51
+ "use_stemmer": self.use_stemmer,
52
+ },
53
+ )
54
+ artifact.add_dir(index_name, name=index_name)
55
  artifact.save()
56
+
57
+ @classmethod
58
+ def from_wandb_artifact(cls, index_artifact_address: str):
59
+ if wandb.run:
60
+ artifact = wandb.run.use_artifact(
61
+ index_artifact_address, type="bm25s-index"
62
+ )
63
+ artifact_dir = artifact.download()
64
+ else:
65
+ api = wandb.Api()
66
+ artifact = api.artifact(index_artifact_address)
67
+ artifact_dir = artifact.download()
68
+ index_name = glob(os.path.join(artifact_dir, "*"))[0].split("/")[-1]
69
+ retriever = bm25s.BM25.load(index_name, load_corpus=True)
70
+ metadata = artifact.metadata
71
+ return cls(
72
+ language=metadata["language"],
73
+ use_stemmer=metadata["use_stemmer"],
74
+ retriever=retriever,
75
+ )
76
+
77
+ @weave.op()
78
+ def retrieve(self, query: str, top_k: int = 2):
79
+ query_tokens = bm25s.tokenize(
80
+ query,
81
+ stopwords=LANGUAGE_DICT[self.language],
82
+ stemmer=Stemmer(self.language) if self.use_stemmer else None,
83
+ )
84
+ results, scores = self._retriever.retrieve(query_tokens, k=top_k)
85
+ return {
86
+ "results": results,
87
+ "scores": scores,
88
+ }