Maksym Batiuk commited on
Commit
e66a368
1 Parent(s): 748993b

add caching

Browse files
Files changed (1) hide show
  1. src/rag/retriever.py +2 -2
src/rag/retriever.py CHANGED
@@ -30,13 +30,13 @@ class Retriever:
30
  ])
31
 
32
  # Load SentenceTransformer model for semantic similarity
33
- self.semantic_model = SentenceTransformer(semantic_model)
34
  self.semantic_chunk_embeddings = self.semantic_model.encode(
35
  self.text_chunks, convert_to_tensor=False
36
  )
37
 
38
  # Load CrossEncoder model for reranking
39
- self.reranker = CrossEncoder(reranker_model)
40
 
41
 
42
  def _load_chunks_from_directory(self, chunked_dir: str) -> list[str]:
 
30
  ])
31
 
32
  # Load SentenceTransformer model for semantic similarity
33
+ self.semantic_model = SentenceTransformer(semantic_model, cache_folder="cache/")
34
  self.semantic_chunk_embeddings = self.semantic_model.encode(
35
  self.text_chunks, convert_to_tensor=False
36
  )
37
 
38
  # Load CrossEncoder model for reranking
39
+ self.reranker = CrossEncoder(reranker_model, cache_dir="cache/")
40
 
41
 
42
  def _load_chunks_from_directory(self, chunked_dir: str) -> list[str]: