Revert "make cpu tensors the default for both testing and production"

This reverts commit 020e1b3.

Author: commit111

app/rag_system.py

@@ -27,16 +27,14 @@ def embed_knowledge_base(self):
     def normalize_query(self, query):
         return query.lower().strip()
 
-    def get_query_embedding(self, query, use_cpu=True):
+    def get_query_embedding(self, query, use_cpu=False):
         normalized_query = self.normalize_query(query)
         query_embedding = self.model.encode([normalized_query], convert_to_tensor=True)
-        # Move the embeddings to the CPU to ensure compatibility with operations like cosine_similarity
         if use_cpu:
             query_embedding = query_embedding.cpu()
         return query_embedding
 
-    def get_doc_embeddings(self, use_cpu=True):
-       # Move the embeddings to the CPU to ensure compatibility with operations like cosine_similarity
+    def get_doc_embeddings(self, use_cpu=False):
         if use_cpu:
             return self.doc_embeddings.cpu()
         return self.doc_embeddings
@@ -64,9 +62,12 @@ def compute_document_scores(self, query_embedding, doc_embeddings, high_match_th
 
         return result
 
-    def retrieve(self, query, similarity_threshold=0.7, high_match_threshold=0.8, max_docs=5):
-        query_embedding = self.get_query_embedding(query)
-        doc_embeddings = self.get_doc_embeddings()
+    def retrieve(self, query, similarity_threshold=0.7, high_match_threshold=0.8, max_docs=5, use_cpu=False):
+        # Note: Set use_cpu=True to run on CPU, which is useful for testing or environments without a GPU.
+        # Set use_cpu=False to leverage GPU for better performance in production.
+        
+        query_embedding = self.get_query_embedding(query, use_cpu)
+        doc_embeddings = self.get_doc_embeddings(use_cpu)
 
         doc_scores = self.compute_document_scores(query_embedding, doc_embeddings, high_match_threshold)
         retrieved_docs = self.get_top_docs(doc_scores, similarity_threshold, max_docs)

app/test_rag_system.py

@@ -53,7 +53,8 @@ def test_get_doc_embeddings(self):
     def test_retrieve_fallback(self):
         # test a query that should return the fallback response
         query = "Hello"
-        result = self.rag_system.retrieve(query)
+        # set use_cpu to True, as testing has no GPU calculations
+        result = self.rag_system.retrieve(query, use_cpu=True)
         self.assertIsInstance(result, list)
         self.assertGreater(len(result), 0)
         self.assertEqual(len(result), 1)  # should return one result for fallback
@@ -66,7 +67,8 @@ def test_retrieve_fallback(self):
     def test_retrieve_actual_response(self):
         # test a query that should return an actual response from the knowledge base
         query = "What is Defang?"
-        result = self.rag_system.retrieve(query)
+        # set use_cpu to True, as testing has no GPU calculations
+        result = self.rag_system.retrieve(query, use_cpu=True)
         self.assertIsInstance(result, list)
         self.assertGreater(len(result), 0)
         self.assertLessEqual(len(result), 5)  # should return up to max_docs (5)
@@ -79,8 +81,8 @@ def test_retrieve_actual_response(self):
     def test_compute_document_scores(self):
         query = "Does Defang have an MCP sample?"
         # get embeddings and move them to CPU, as testing has no GPU calculations
-        query_embedding = self.rag_system.get_query_embedding(query)
-        doc_embeddings = self.rag_system.get_doc_embeddings()
+        query_embedding = self.rag_system.get_query_embedding(query, use_cpu=True)
+        doc_embeddings = self.rag_system.get_doc_embeddings(use_cpu=True)
 
         # call function and get results
         result = self.rag_system.compute_document_scores(query_embedding, doc_embeddings, high_match_threshold=0.8)