remove use_cpu args

as we will now always run on cpu--these are irrelevant.

Author: jordanstephens

app/rag_system.py

@@ -30,17 +30,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):
         normalized_query = self.normalize_query(query)
         query_embedding = self.model.encode([normalized_query], convert_to_tensor=True)
-        if use_cpu:
-            query_embedding = query_embedding.cpu()
+        query_embedding = query_embedding.cpu()
         return query_embedding
 
-    def get_doc_embeddings(self, use_cpu=True):
-        if use_cpu:
-            return self.doc_embeddings.cpu()
-        return self.doc_embeddings
+    def get_doc_embeddings(self):
+        return self.doc_embeddings.cpu()
 
     def compute_document_scores(self, query_embedding, doc_embeddings, high_match_threshold):
         text_similarities = cosine_similarity(query_embedding, doc_embeddings)[0]
@@ -66,12 +63,9 @@ def compute_document_scores(self, query_embedding, doc_embeddings, high_match_th
 
         return result
 
-    def retrieve(self, query, similarity_threshold=0.4, high_match_threshold=0.8, max_docs=5, use_cpu=True):
-        # 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)
+    def retrieve(self, query, similarity_threshold=0.4, high_match_threshold=0.8, max_docs=5):
+        query_embedding = self.get_query_embedding(query)
+        doc_embeddings = self.get_doc_embeddings()
 
         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,8 +53,7 @@ def test_get_doc_embeddings(self):
     def test_retrieve_fallback(self):
         # test a query that should return the fallback response
         query = "Hello"
-        # set use_cpu to True, as testing has no GPU calculations
-        result = self.rag_system.retrieve(query, use_cpu=True)
+        result = self.rag_system.retrieve(query)
         self.assertIsInstance(result, list)
         self.assertGreater(len(result), 0)
         self.assertEqual(len(result), 1)  # should return one result for fallback
@@ -67,8 +66,7 @@ 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?"
-        # set use_cpu to True, as testing has no GPU calculations
-        result = self.rag_system.retrieve(query, use_cpu=True)
+        result = self.rag_system.retrieve(query)
         self.assertIsInstance(result, list)
         self.assertGreater(len(result), 0)
         self.assertLessEqual(len(result), 5)  # should return up to max_docs (5)
@@ -80,9 +78,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, use_cpu=True)
-        doc_embeddings = self.rag_system.get_doc_embeddings(use_cpu=True)
+        query_embedding = self.rag_system.get_query_embedding(query)
+        doc_embeddings = self.rag_system.get_doc_embeddings()
 
         # call function and get results
         result = self.rag_system.compute_document_scores(query_embedding, doc_embeddings, high_match_threshold=0.8)
@@ -105,4 +102,4 @@ def test_compute_document_scores(self):
         print("Test for compute_document_scores passed successfully!")
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()