Upload retriever.py

rag_101/retriever.py

@@ -114,7 +114,7 @@ def retrieve_context(query, retriever, reranker_model):
 
 
 def load_embedding_model(
-    model_name: str = "BAAI/bge-large-en-v1.5", device: str = "CPU"
+    model_name: str = "BAAI/bge-large-en-v1.5", device: str = "cpu"
 ) -> HuggingFaceEmbeddings:
     model_kwargs = {"device": device}
     encode_kwargs = {
@@ -129,7 +129,7 @@ def load_embedding_model(
 
 
 def load_reranker_model(
-    reranker_model_name: str = "BAAI/bge-reranker-large", device: str = "CPU"
+    reranker_model_name: str = "BAAI/bge-reranker-large", device: str = "cpu"
 ) -> CrossEncoder:
     reranker_model = CrossEncoder(
         model_name=reranker_model_name, max_length=1024, device=device