Maksym Batiuk
commited on
Commit
•
e66a368
1
Parent(s):
748993b
add caching
Browse files- 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]:
|