Update policy_retriever_2.py

policy_retriever_2.py

@@ -52,7 +52,7 @@ class PolicyRetriever:
             outputs = self.model(**inputs)
         return outputs.last_hidden_state[:, 0, :].cpu().numpy()[0].tolist()
     
-   def search(self, query: str, top_k: int = 5) -> List[Dict]:
+    def search(self, query: str, top_k: int = 5) -> List[Dict]:
         """Search for relevant policy sections based on query."""
         print("\nSearching for relevant policy sections...")
         query_embedding = self.get_query_embedding(query)
@@ -113,14 +113,9 @@ Policy sections:
         
         # Add chunks to prompt
         for i, chunk in enumerate(chunks, 1):
-            prompt += f"
-Section {i}:
-Heading: {chunk['heading']}
-Content: {chunk['content']}
-"
+            prompt += f"\nSection {i}:\nHeading: {chunk['heading']}\nContent: {chunk['content']}\n"
         
-        prompt += "
-Provide a clear, direct answer based only on the information shown above:"
+        prompt += "\nProvide a clear, direct answer based only on the information shown above:"
 
         # Generate answer
         inputs = self.llm_tokenizer(
@@ -145,11 +140,9 @@ Provide a clear, direct answer based only on the information shown above:"
         answer = self.llm_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
         
         # Format and print response
-        print(f"
-Q: {query}")
+        print(f"\nQ: {query}")
         print(f"A: {answer}")
-        print("
-Based on sections:")
+        print("\nBased on sections:")
         for chunk in chunks:
             print(f"- {chunk['heading']}")
         
@@ -157,15 +150,13 @@ Based on sections:")
 
     def search_and_generate(self, query: str, top_k: int = 5) -> str:
         """Combined search and answer generation."""
-        print(f"
-Processing query: {query}")
+        print(f"\nProcessing query: {query}")
         retrieved_chunks = self.search(query, top_k=top_k)
         return self.generate_answer(query, retrieved_chunks)
 
 def test_retriever():
     """Test function for the PolicyRetriever."""
-    print("
-=== Testing Policy Retriever ===")
+    print("\n=== Testing Policy Retriever ===")
     retriever = PolicyRetriever()
 
     # Standard test queries
@@ -179,9 +170,7 @@ def test_retriever():
     # Run tests
     for query in test_queries:
         answer = retriever.search_and_generate(query)
-        print("
-" + "-"*70 + "
-")
+        print("\n" + "-"*70 + "\n")
 
 if __name__ == "__main__":
-    test_retriever()
+    test_retriever()