Migrate to Llama.cpp

app/rag.py

@@ -1,59 +1,100 @@
-import os
-import logging
-
 from llama_index.core import (
     SimpleDirectoryReader,
-    VectorStoreIndex,
+    # VectorStoreIndex,
     StorageContext,
     Settings,
     get_response_synthesizer)
 from llama_index.core.query_engine import RetrieverQueryEngine, TransformQueryEngine
 from llama_index.core.node_parser import SentenceSplitter
 from llama_index.core.schema import TextNode, MetadataMode
-from llama_index.vector_stores.qdrant import QdrantVectorStore
-from llama_index.embeddings.ollama import OllamaEmbedding
-from llama_index.llms.ollama import Ollama
 from llama_index.core.retrievers import VectorIndexRetriever
-from llama_index.core.indices.query.query_transform import HyDEQueryTransform
+# from llama_index.core.indices.query.query_transform import HyDEQueryTransform
+
+from llama_index.core.response_synthesizers import ResponseMode
+# from transformers import AutoTokenizer
+from llama_index.core.vector_stores import VectorStoreQuery
+from llama_index.core.indices.vector_store.base import VectorStoreIndex
+from llama_index.vector_stores.qdrant import QdrantVectorStore
 from qdrant_client import QdrantClient
+import logging
+
+from llama_index.llms.llama_cpp import LlamaCPP
+from llama_index.embeddings.fastembed import FastEmbedEmbedding
 
-QDRANT_API_URL = os.getenv('QDRANT_API_URL')
-QDRANT_API_KEY = os.getenv('QDRANT_API_KEY')
 
 
 class ChatPDF:
-    hyde_query_engine = None
-    text_parser = None
-    vector_store = None
-    embed_model = None
-    logger = None
+    logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger(__name__)
+    query_engine = None
 
-    def __init__(self):
-        logging.basicConfig(level=logging.INFO)
-        self.logger = logging.getLogger(__name__)
+    # 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 = ""
+        for message in messages:
+            if message.role == 'system':
+                prompt += f"<|system|>\n{message.content}</s>\n"
+            elif message.role == 'user':
+                prompt += f"<|user|>\n{message.content}</s>\n"
+            elif message.role == 'assistant':
+                prompt += f"<|assistant|>\n{message.content}</s>\n"
+
+        if not prompt.startswith("<|system|>\n"):
+            prompt = "<|system|>\n</s>\n" + prompt
+
+        prompt = prompt + "<|assistant|>\n"
+
+        return prompt
 
-        self.text_parser = SentenceSplitter(chunk_size=512, chunk_overlap=100)
+    def completion_to_prompt(completion):
+        return f"<|system|>\n</s>\n<|user|>\n{completion}</s>\n<|assistant|>\n"
+
+
+    def __init__(self):
+        text_parser = SentenceSplitter(chunk_size=512, chunk_overlap=20)
 
         self.logger.info("initializing the vector store related objects")
-        client = QdrantClient(url=QDRANT_API_URL, api_key=QDRANT_API_KEY)
+        # client = QdrantClient(host="localhost", port=6333)
+        client = QdrantClient(":memory:")
         self.vector_store = QdrantVectorStore(client=client, collection_name="rag_documents")
 
-        self.logger.info("initializing the OllamaEmbedding")
-        self.embed_model = OllamaEmbedding(model_name='mxbai-embed-large')
+        self.logger.info("initializing the FastEmbedEmbedding")
+        self.embed_model = FastEmbedEmbedding(model_name="BAAI/bge-small-en")
+
+        llm = LlamaCPP(
+            # model_url=self.model_url,
+            temperature=0.1,
+            max_new_tokens=256,
+            context_window=3900,
+            # generate_kwargs={},
+            model_kwargs={"n_gpu_layers": -1},
+            messages_to_prompt=self.messages_to_prompt,
+            completion_to_prompt=self.completion_to_prompt,
+            verbose=True,
+        )
+
+        # tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct")
+        # tokenizer.save_pretrained("./models/tokenizer/")
+
         self.logger.info("initializing the global settings")
+        Settings.text_splitter = text_parser
         Settings.embed_model = self.embed_model
-        Settings.llm = Ollama(model="qwen:1.8b", request_timeout=1000000)
-        Settings.transformations = [self.text_parser]
+        Settings.llm = llm
+        # Settings.tokenzier = tokenizer
+        Settings.transformations = [text_parser]
 
-    def ingest(self, dir_path: str):
-        docs = SimpleDirectoryReader(input_dir=dir_path).load_data()
+    def ingest(self, pdf_file_path: str):
         text_chunks = []
         doc_ids = []
         nodes = []
 
+        docs = SimpleDirectoryReader(input_dir="files").load_data()
+
         self.logger.info("enumerating docs")
         for doc_idx, doc in enumerate(docs):
-            curr_text_chunks = self.text_parser.split_text(doc.text)
+            curr_text_chunks = text_parser.split_text(doc.text)
             text_chunks.extend(curr_text_chunks)
             doc_ids.extend([doc_idx] * len(curr_text_chunks))
 
@@ -80,26 +121,37 @@ class ChatPDF:
             transformations=Settings.transformations,
         )
 
-        self.logger.info("initializing the VectorIndexRetriever with top_k as 5")
-        vector_retriever = VectorIndexRetriever(index=index, similarity_top_k=5)
-        response_synthesizer = get_response_synthesizer()
-        self.logger.info("creating the RetrieverQueryEngine instance")
-        vector_query_engine = RetrieverQueryEngine(
-            retriever=vector_retriever,
+        self.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,
+        )
+
+        self.logger.info("assemble query engine")
+        self.query_engine = RetrieverQueryEngine(
+            retriever=retriever,
             response_synthesizer=response_synthesizer,
         )
-        self.logger.info("creating the HyDEQueryTransform instance")
-        hyde = HyDEQueryTransform(include_original=True)
-        self.hyde_query_engine = TransformQueryEngine(vector_query_engine, hyde)
+
+        # self.logger.info("creating the HyDEQueryTransform instance")
+        # hyde = HyDEQueryTransform(include_original=True)
+        # self.hyde_query_engine = TransformQueryEngine(vector_query_engine, hyde)
 
     def ask(self, query: str):
-        if not self.hyde_query_engine:
+        if not self.query_engine:
             return "Please, add a PDF document first."
 
         self.logger.info("retrieving the response to the query")
-        response = self.hyde_query_engine.query(str_or_query_bundle=query)
-        self.logger.info(response)
+        response = self.query_engine.query(str_or_query_bundle=query)
+        print(response)
         return response
 
     def clear(self):
-        self.hyde_query_engine = None
+        self.query_engine = None

requirements.txt

@@ -1,5 +1,8 @@
 fastapi
 llama-index
 llama-index-vector-stores-qdrant
-llama-index-embeddings-ollama
-llama-index-llms-ollama
+qdrant-client
+python-dotenv
+llama-index-llms-llama-cpp
+llama-index-embeddings-fastembed
+fastembed==0.2.7