anpigon commited on
Commit
8eff983
·
1 Parent(s): c80fa93

refactor: Update retrievers.py to improve retriever configuration and naming conventions

Browse files
Files changed (1) hide show
  1. 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(run_name="faiss")
 
 
28
 
29
- bm25_retriever = load_bm25_retriever().with_config(run_name="bm25")
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