Update app/rag.py

app/rag.py

@@ -4,8 +4,7 @@ from llama_index.core import (
     SimpleDirectoryReader,
     VectorStoreIndex,
     StorageContext,
-    Settings,
-    get_response_synthesizer)
+    Settings)
 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.core.schema import TextNode, MetadataMode
 from llama_index.core.vector_stores import VectorStoreQuery
@@ -20,6 +19,8 @@ store_dir = os.path.expanduser("~/wtp_be_store/")
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+model_url = "https://huggingface.co/Qwen/Qwen2-0.5B-Instruct-GGUF/resolve/main/qwen2-0_5b-instruct-q4_k_m.gguf"
+
 class ChatPDF:
     pdf_count = 0
     text_chunks = []
@@ -33,33 +34,26 @@ class ChatPDF:
         self.client = QdrantClient(path=store_dir)
         self.vector_store = QdrantVectorStore(
             client=self.client,
-            collection_name="rag_documents",
-            # enable_hybrid=True
+            collection_name="rag_documents"
         )
 
         logger.info("initializing the FastEmbedEmbedding")
-        self.embed_model = FastEmbedEmbedding(
-            # model_name="BAAI/bge-small-en"
-        )
+        self.embed_model = FastEmbedEmbedding()
 
         llm = LlamaCPP(
-            model_url="https://huggingface.co/Qwen/Qwen2-0.5B-Instruct-GGUF/resolve/main/qwen2-0_5b-instruct-fp16.gguf",
+            model_url=model_url,
             temperature=0.1,
+            model_path=None,
             max_new_tokens=256,
-            generate_kwargs={"max_tokens": 256, "temperature": 0.1, "top_k": 3},
-            # messages_to_prompt=self.messages_to_prompt,
-            # completion_to_prompt=self.completion_to_prompt,
+            context_window=29440,
+            generate_kwargs={},
             verbose=True,
         )
 
-        # tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
-        # tokenizer.save_pretrained("./models/tokenizer/")
-
         logger.info("initializing the global settings")
         Settings.text_splitter = self.text_parser
         Settings.embed_model = self.embed_model
         Settings.llm = llm
-        # Settings.tokenzier = tokenizer
         Settings.transformations = [self.text_parser]
 
     def ingest(self, files_dir: str):
@@ -73,10 +67,8 @@ class ChatPDF:
             self.doc_ids.extend([doc_idx] * len(curr_text_chunks))
 
         logger.info("enumerating text_chunks")
-        for idx, text_chunk in enumerate(self.text_chunks):
+        for text_chunk in self.text_chunks:
             node = TextNode(text=text_chunk)
-            # src_doc = docs[self.doc_ids[idx]]
-            # node.metadata = src_doc.metadata
             if node.get_content(metadata_mode=MetadataMode.EMBED):
                 self.nodes.append(node)
 
@@ -93,7 +85,7 @@ class ChatPDF:
         index = VectorStoreIndex(
             nodes=self.nodes,
             storage_context=storage_context,
-            transformations=Settings.transformations,
+            transformations=Settings.transformations
         )
 
         self.query_engine = index.as_query_engine(
@@ -103,14 +95,13 @@ class ChatPDF:
 
     def ask(self, query: str):
         logger.info("retrieving the response to the query")
-        streaming_response = self.query_engine.query(query)
+        streaming_response = self.query_engine.query("You are an assistant for question-answering tasks. Use three \
+            sentences only and keep the answer concise.\n\n" + query)
         return streaming_response
 
     def clear(self):
-        # self.vector_store.clear()
-        if self.nodes:
-            self.vector_store.delete_nodes(self.nodes)
-            self.pdf_count = 0
-            self.text_chunks = []
-            self.doc_ids = []
-            self.nodes = []
+        self.vector_store.clear()
+        self.pdf_count = 0
+        self.text_chunks = []
+        self.doc_ids = []
+        self.nodes = []