Update rag.py

app/rag.py

@@ -27,13 +27,14 @@ from llama_index.embeddings.fastembed import FastEmbedEmbedding
 QDRANT_API_URL = os.getenv('QDRANT_API_URL')
 QDRANT_API_KEY = os.getenv('QDRANT_API_KEY')
 
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 class ChatPDF:
-    logging.basicConfig(level=logging.INFO)
-    logger = logging.getLogger(__name__)
     query_engine = None
 
-    # model_url = "https://huggingface.co/Qwen/Qwen1.5-1.8B-Chat-GGUF/resolve/main/qwen1_5-1_8b-chat-q8_0.gguf"
-    model_url = "https://huggingface.co/microsoft/Phi-3-mini-4k-instruct-gguf/resolve/main/Phi-3-mini-4k-instruct-q4.gguf"
+    model_url = "https://huggingface.co/Qwen/Qwen1.5-1.8B-Chat-GGUF/resolve/main/qwen1_5-1_8b-chat-q8_0.gguf"
+    # model_url = "https://huggingface.co/microsoft/Phi-3-mini-4k-instruct-gguf/resolve/main/Phi-3-mini-4k-instruct-q4.gguf"
 
     # def messages_to_prompt(messages):
     #     prompt = ""
@@ -59,7 +60,7 @@ class ChatPDF:
     def __init__(self):
         self.text_parser = SentenceSplitter(chunk_size=512, chunk_overlap=20)
 
-        self.logger.info("initializing the vector store related objects")
+        logger.info("initializing the vector store related objects")
         # client = QdrantClient(host="localhost", port=6333)
         # client = QdrantClient(url=QDRANT_API_URL, api_key=QDRANT_API_KEY)
         client = QdrantClient(":memory:")
@@ -69,7 +70,7 @@ class ChatPDF:
             # enable_hybrid=True
         )
 
-        self.logger.info("initializing the FastEmbedEmbedding")
+        logger.info("initializing the FastEmbedEmbedding")
         self.embed_model = FastEmbedEmbedding(
             # model_name="BAAI/bge-small-en"
         )
@@ -89,7 +90,7 @@ class ChatPDF:
         # tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
         # tokenizer.save_pretrained("./models/tokenizer/")
 
-        self.logger.info("initializing the global settings")
+        logger.info("initializing the global settings")
         Settings.text_splitter = self.text_parser
         Settings.embed_model = self.embed_model
         Settings.llm = llm
@@ -103,55 +104,57 @@ class ChatPDF:
 
         docs = SimpleDirectoryReader(input_dir=files_dir).load_data()
 
-        self.logger.info("enumerating docs")
+        logger.info("enumerating docs")
         for doc_idx, doc in enumerate(docs):
             curr_text_chunks = self.text_parser.split_text(doc.text)
             text_chunks.extend(curr_text_chunks)
             doc_ids.extend([doc_idx] * len(curr_text_chunks))
 
-        self.logger.info("enumerating text_chunks")
+        logger.info("enumerating text_chunks")
         for idx, text_chunk in enumerate(text_chunks):
             node = TextNode(text=text_chunk)
             src_doc = docs[doc_ids[idx]]
             node.metadata = src_doc.metadata
             nodes.append(node)
 
-        self.logger.info("enumerating nodes")
+        logger.info("enumerating nodes")
         for node in nodes:
             node_embedding = self.embed_model.get_text_embedding(
                 node.get_content(metadata_mode=MetadataMode.ALL)
             )
             node.embedding = node_embedding
 
-        self.logger.info("initializing the storage context")
+        logger.info("initializing the storage context")
         storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
-        self.logger.info("indexing the nodes in VectorStoreIndex")
+        logger.info("indexing the nodes in VectorStoreIndex")
         index = VectorStoreIndex(
             nodes=nodes,
             storage_context=storage_context,
             transformations=Settings.transformations,
         )
 
-        self.logger.info("configure retriever")
-        retriever = VectorIndexRetriever(
-            index=index,
-            similarity_top_k=6,
-            # vector_store_query_mode="hybrid"
-        )
+        # logger.info("configure retriever")
+        # retriever = VectorIndexRetriever(
+        #     index=index,
+        #     similarity_top_k=6,
+        #     # vector_store_query_mode="hybrid"
+        # )
 
-        self.logger.info("configure response synthesizer")
-        response_synthesizer = get_response_synthesizer(
-            # streaming=True,
-            response_mode=ResponseMode.COMPACT,
-        )
+        # logger.info("configure response synthesizer")
+        # response_synthesizer = get_response_synthesizer(
+        #     # streaming=True,
+        #     response_mode=ResponseMode.COMPACT,
+        # )
 
-        self.logger.info("assemble query engine")
-        self.query_engine = RetrieverQueryEngine(
-            retriever=retriever,
-            response_synthesizer=response_synthesizer,
-        )
+        # logger.info("assemble query engine")
+        # self.query_engine = RetrieverQueryEngine(
+        #     retriever=retriever,
+        #     response_synthesizer=response_synthesizer,
+        # )
+
+        self.query_engine = index.as_query_engine()
 
-        # self.logger.info("creating the HyDEQueryTransform instance")
+        # logger.info("creating the HyDEQueryTransform instance")
         # hyde = HyDEQueryTransform(include_original=True)
         # self.hyde_query_engine = TransformQueryEngine(vector_query_engine, hyde)
 
@@ -159,8 +162,9 @@ class ChatPDF:
         if not self.query_engine:
             return "Please, add a PDF document first."
 
-        self.logger.info("retrieving the response to the query")
-        response = self.query_engine.query(str_or_query_bundle=query)
+        logger.info("retrieving the response to the query")
+        # response = self.query_engine.query(str_or_query_bundle=query)
+        response = self.query_engine.query(query)
         print(response)
         return response