Spaces:
Sleeping
Sleeping
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(
|
|
|
|
|
42 |
if wandb.run:
|
43 |
-
artifact = wandb.Artifact(
|
44 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
+
}
|