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
|