refactor: Update retrievers.py to improve retriever configuration and naming conventions
Browse files- libs/retrievers.py +7 -4
libs/retrievers.py
CHANGED
@@ -24,19 +24,22 @@ def load_faiss_retriever(embeddings):
|
|
24 |
|
25 |
|
26 |
def load_retrievers(embeddings):
|
27 |
-
faiss_retriever = load_faiss_retriever(embeddings).with_config(
|
|
|
|
|
28 |
|
29 |
-
bm25_retriever = load_bm25_retriever().with_config(run_name="
|
30 |
|
31 |
ensemble_retriever = EnsembleRetriever(
|
32 |
retrievers=[bm25_retriever, faiss_retriever],
|
33 |
weights=[0.7, 0.3],
|
34 |
search_type="mmr",
|
35 |
-
)
|
36 |
|
37 |
compressor = CohereRerank(model="rerank-multilingual-v3.0", top_n=5)
|
38 |
compression_retriever = ContextualCompressionRetriever(
|
39 |
base_compressor=compressor,
|
40 |
base_retriever=ensemble_retriever,
|
41 |
-
)
|
|
|
42 |
return compression_retriever
|
|
|
24 |
|
25 |
|
26 |
def load_retrievers(embeddings):
|
27 |
+
faiss_retriever = load_faiss_retriever(embeddings).with_config(
|
28 |
+
run_name="FaissRetriever"
|
29 |
+
)
|
30 |
|
31 |
+
bm25_retriever = load_bm25_retriever().with_config(run_name="BM25Retriever")
|
32 |
|
33 |
ensemble_retriever = EnsembleRetriever(
|
34 |
retrievers=[bm25_retriever, faiss_retriever],
|
35 |
weights=[0.7, 0.3],
|
36 |
search_type="mmr",
|
37 |
+
).with_config(run_name="EnsembleRetriever")
|
38 |
|
39 |
compressor = CohereRerank(model="rerank-multilingual-v3.0", top_n=5)
|
40 |
compression_retriever = ContextualCompressionRetriever(
|
41 |
base_compressor=compressor,
|
42 |
base_retriever=ensemble_retriever,
|
43 |
+
).with_config(run_name="ContextualCompressionRetriever")
|
44 |
+
|
45 |
return compression_retriever
|