99 additions of files in the repo, 99 additions of files...

FAISS-index.py

@@ -1,27 +1,25 @@
 from datasets import Dataset, load_from_disk
 import faiss
 import numpy as np
-from transformers import RagRetriever, RagTokenizer, RagSequenceForGeneration
+from transformers import RagTokenizer, RagSequenceForGeneration
 
-# Example: Create a dataset
-data = {"text": ["This is a sample text.", "Another sample text."]}
-dataset = Dataset.from_dict(data)
+def create_and_save_faiss_index(dataset_path, index_path):
+    dataset = load_from_disk(dataset_path)
+    passages = dataset["text"]
+    
+    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
+    model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")
+    
+    passage_embeddings = model.get_encoder()(
+        tokenizer(passages, return_tensors="pt", padding=True, truncation=True)
+    ).last_hidden_state.mean(dim=1).detach().numpy()
 
-# Save the dataset to disk
-dataset_path = "path/to/your/dataset"
-dataset.save_to_disk(dataset_path)
+    index = faiss.IndexFlatL2(passage_embeddings.shape[1])
+    index.add(passage_embeddings)
 
-# Create FAISS index
-passages = dataset["text"]
-tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
-model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")
-passage_embeddings = model.get_encoder()(
-    tokenizer(passages, return_tensors="pt", padding=True, truncation=True)
-).last_hidden_state.mean(dim=1).detach().numpy()
+    faiss.write_index(index, index_path)
 
-index = faiss.IndexFlatL2(passage_embeddings.shape[1])
-index.add(passage_embeddings)
-
-# Save the index to disk
-index_path = "path/to/your/index"
-faiss.write_index(index, index_path)
+if __name__ == "__main__":
+    dataset_path = "path/to/your/hf_dataset"
+    index_path = "path/to/your/hf_index"
+    create_and_save_faiss_index(dataset_path, index_path)

RAGbot.py

@@ -10,7 +10,12 @@ from langchain.document_loaders import PyPDFLoader
 from langchain.prompts import PromptTemplate
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import spaces
-from langchain_text_splitters import CharacterTextSplitter, RecursiveCharacterTextSplitter
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType
+from datasets import Dataset, load_from_disk
+import faiss
+import numpy as np
+from pastebin_api import get_protected_content
 
 class RAGbot:
     def __init__(self, config_path="config.yaml"):
@@ -20,7 +25,8 @@ class RAGbot:
         self.prompt = None
         self.documents = None
         self.embeddings = None
-        self.vectordb = None
+        self.zilliz_vectordb = None
+        self.hf_vectordb = None
         self.tokenizer = None
         self.model = None
         self.pipeline = None
@@ -38,22 +44,26 @@ class RAGbot:
             self.model_embeddings = config["modelEmbeddings"]
             self.auto_tokenizer = config["autoTokenizer"]
             self.auto_model_for_causal_lm = config["autoModelForCausalLM"]
+            self.zilliz_config = config["zilliz"]
+            self.persona_paste_key = config["personaPasteKey"]
+
+    def connect_to_zilliz(self):
+        connections.connect(
+            host=self.zilliz_config["host"],
+            port=self.zilliz_config["port"],
+            user=self.zilliz_config["user"],
+            password=self.zilliz_config["password"],
+            secure=True
+        )
+        self.zilliz_vectordb = Collection(self.zilliz_config["collection"])
 
     def load_embeddings(self):
         self.embeddings = HuggingFaceEmbeddings(model_name=self.model_embeddings)
-        print("Embedding model loaded")
-
-    def load_vectordb(self):
-        overlap = int((self.overlap_percentage / 100) * self.chunk_size)
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=self.chunk_size,
-            chunk_overlap=overlap,
-            length_function=len,
-            add_start_index=True,
-        )
-        docs = text_splitter.split_documents(self.documents)
-        self.vectordb = Chroma.from_documents(docs, self.embeddings)
-        print("Vector store created")
+
+    def load_hf_vectordb(self, dataset_path, index_path):
+        dataset = load_from_disk(dataset_path)
+        index = faiss.read_index(index_path)
+        self.hf_vectordb = (dataset, index)
 
     @spaces.GPU
     def load_tokenizer(self):
@@ -67,20 +77,34 @@ class RAGbot:
             model_kwargs={"torch_dtype": torch.bfloat16},
             device="cuda",
         )
-        print("Model pipeline loaded")
 
-    def get_organic_context(self, query):
-        documents = self.vectordb.similarity_search_with_relevance_scores(query, k=self.max_chunks_in_context)
-        context = self.format_seperator.join([doc.page_content for doc, score in documents])
+    def get_organic_context(self, query, use_hf=False):
+        if use_hf:
+            dataset, index = self.hf_vectordb
+            D, I = index.search(np.array([self.embeddings.embed_query(query)]), self.max_chunks_in_context)
+            context = self.format_seperator.join([dataset[i] for i in I[0]])
+        else:
+            result = self.zilliz_vectordb.search(
+                data=[self.embeddings.embed_query(query)],
+                anns_field="embeddings",
+                param={"metric_type": "IP", "params": {"nprobe": 10}},
+                limit=self.max_chunks_in_context,
+                expr=None,
+            )
+            context = self.format_seperator.join([hit.entity.get('text') for hit in result[0]])
+        
         self.current_context = context
-        print("Context Ready")
-        print(self.current_context)
+
+    def load_persona_data(self):
+        persona_content = get_protected_content(self.persona_paste_key)
+        persona_data = yaml.safe_load(persona_content)
+        self.persona_text = persona_data["persona_text"]
 
     @spaces.GPU
-    def create_organic_response(self, history, query):
-        self.get_organic_context(query)
+    def create_organic_response(self, history, query, use_hf=False):
+        self.get_organic_context(query, use_hf=use_hf)
         messages = [
-            {"role": "system", "content": "From the context given below, answer the user's question\n" + self.current_context},
+            {"role": "system", "content": f"Based on the given context, answer the user's question while maintaining the persona:\n{self.persona_text}\n\nContext:\n{self.current_context}"},
             {"role": "user", "content": query},
         ]
 
@@ -97,17 +121,15 @@ class RAGbot:
             temperature=temp,
             top_p=0.9,
         )
-        print(outputs)
         return outputs[0]["generated_text"][len(prompt):]
 
     def process_file(self, file):
         self.documents = PyPDFLoader(file.name).load()
         self.load_embeddings()
-        self.load_vectordb()
-        self.create_organic_pipeline()
+        self.connect_to_zilliz()
 
     @spaces.GPU
-    def generate_response(self, history, query, file, chunk_size, chunk_overlap_percentage, model_temperature, max_chunks_in_context):
+    def generate_response(self, history, query, file, chunk_size, chunk_overlap_percentage, model_temperature, max_chunks_in_context, use_hf_index=False, hf_dataset_path=None, hf_index_path=None):
         self.chunk_size = chunk_size
         self.overlap_percentage = chunk_overlap_percentage
         self.model_temperatue = model_temperature
@@ -115,19 +137,28 @@ class RAGbot:
 
         if not query:
             raise gr.Error(message='Submit a question')
-        if not file:
-            raise gr.Error(message='Upload a PDF')
-        if not self.processed:
-            self.process_file(file)
-            self.processed = True
 
-        result = self.create_organic_response(history="", query=query)
+        if use_hf_index:
+            if not hf_dataset_path or not hf_index_path:
+                raise gr.Error(message='Provide HuggingFace dataset and index paths')
+            self.load_hf_vectordb(hf_dataset_path, hf_index_path)
+            result = self.create_organic_response(history="", query=query, use_hf=True)
+        else:
+            if not file:
+                raise gr.Error(message='Upload a PDF')
+            if not self.processed:
+                self.process_file(file)
+                self.processed = True
+            result = self.create_organic_response(history="", query=query)
+
+        self.load_persona_data()
+        result = f"{self.persona_text}\n\n{result}"
+        
         for char in result:
             history[-1][-1] += char
         return history, ""
 
     def render_file(self, file, chunk_size, chunk_overlap_percentage, model_temperature, max_chunks_in_context):
-        print(chunk_size)
         doc = fitz.open(file.name)
         page = doc[self.page]
         self.chunk_size = chunk_size
@@ -142,4 +173,4 @@ class RAGbot:
         if not text:
             raise gr.Error('Enter text')
         history.append((text, ''))
-        return history
+        return history

requirements.txt

@@ -6,7 +6,15 @@ langchain-community
 tqdm
 accelerate
 pypdf
+faiss-cpu
 protobuf>=3.20,<5
 poetry
+pymilvus
+chromadb
+gradio
+fitz
+PyYAML
+datasets
+numpy
 requests
-pymilvus
+python-dotenv