Update utils/qa.py

utils/qa.py

@@ -1,58 +1,58 @@
-import logging
-from ingestion import DocumentProcessor
-from llm import LLMProcessor
-
-
-class QAEngine:
-    def __init__(self):
-        self.processor = DocumentProcessor()
-        self.llm_processor = LLMProcessor()
-
-    def query(self, question: str, k: int = 5) -> str:
-        """Query the document using semantic search and generate an answer"""
-        query_embedding = self.llm_processor.embed_model.encode(question)
-
-        # Corrected ChromaDB query syntax
-        results = self.processor.index.query(
-            query_embeddings=[query_embedding],
-            n_results=k
-        )
-
-        # Extracting results properly
-        chunks = []
-        for i in range(len(results["documents"][0])):  # Iterate over top-k results
-            chunks.append({
-                "text": results["documents"][0][i],
-                "headings": results["metadatas"][0][i].get("headings", "[]"),
-                "page": results["metadatas"][0][i].get("page"),
-                "content_type": results["metadatas"][0][i].get("content_type")
-            })
-
-        print(f"\nRelevant chunks for query: '{question}'")
-        print("=" * 80)
-
-        context = self.llm_processor.format_context(chunks)
-        print(context)
-
-        return self.llm_processor.generate_answer(context, question)
-
-
-# def main():
-#     logging.basicConfig(level=logging.INFO)
-
-#     processor = DocumentProcessor()
-    
-#     pdf_path = "sample/InternLM.pdf"
-#     processor.process_document(pdf_path)
-
-#     qa_engine = QAEngine()
-#     question = "What are the main features of InternLM-XComposer-2.5?"
-#     answer = qa_engine.query(question)
-
-#     print("\nAnswer:")
-#     print("=" * 80)
-#     print(answer)
-
-
-# if __name__ == "__main__":
-#     main()
+import logging
+from ingestion import DocumentProcessor
+from llm import LLMProcessor
+
+
+class QAEngine:
+    def __init__(self):
+        self.processor = DocumentProcessor()
+        self.llm_processor = LLMProcessor()
+
+    def query(self, question: str, k: int = 5) -> str:
+        """Query the document using semantic search and generate an answer"""
+        query_embedding = self.llm_processor.embed_model.encode(question)
+
+        # Corrected ChromaDB query syntax
+        results = self.processor.index.query(
+            query_embeddings=[query_embedding],
+            n_results=k
+        )
+
+        # Extracting results properly
+        chunks = []
+        for i in range(len(results["documents"][0])):  # Iterate over top-k results
+            chunks.append({
+                "text": results["documents"][0][i],
+                "headings": results["metadatas"][0][i].get("headings", "[]"),
+                "page": results["metadatas"][0][i].get("page"),
+                "content_type": results["metadatas"][0][i].get("content_type")
+            })
+
+        print(f"\nRelevant chunks for query: '{question}'")
+        print("=" * 80)
+
+        context = self.llm_processor.format_context(chunks)
+        print(context)
+
+        return self.llm_processor.generate_answer(context, question)
+
+
+# def main():
+#     logging.basicConfig(level=logging.INFO)
+
+#     processor = DocumentProcessor()
+    
+#     pdf_path = "sample/InternLM.pdf"
+#     processor.process_document(pdf_path)
+
+#     qa_engine = QAEngine()
+#     question = "What are the main features of InternLM-XComposer-2.5?"
+#     answer = qa_engine.query(question)
+
+#     print("\nAnswer:")
+#     print("=" * 80)
+#     print(answer)
+
+
+# if __name__ == "__main__":
+#     main()