Update handler.py

handler.py

@@ -2,7 +2,7 @@ import torch
 import locale
 import os
 from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 from langchain.llms import HuggingFacePipeline
 from langchain.retrievers.document_compressors import LLMChainExtractor
 from langchain.retrievers import ContextualCompressionRetriever
@@ -26,46 +26,30 @@ from langchain.schema import format_document
 from langchain.memory import ConversationBufferMemory
 from langchain_core.messages import AIMessage, HumanMessage, get_buffer_string
 from langchain_core.runnables import RunnableParallel
-    
+from typing import Optional
+from langchain.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores.utils import DistanceStrategy
+
+EMBEDDING_MODEL_NAME = "mixedbread-ai/mxbai-embed-large-v1"
+MARKDOWN_SEPARATORS = [
+    "\n#{1,6} ",
+    "```\n",
+    "\n\\*\\*\\*+\n",
+    "\n---+\n",
+    "\n___+\n",
+    "\n\n",
+    "\n",
+    " ",
+    "",
+]
+
 class EndpointHandler():
-    def split_documents(
-        chunk_size: int,
-        knowledge_base: [],
-        tokenizer_name: Optional[str] = EMBEDDING_MODEL_NAME,
-    ):
-        """
-        Split documents into chunks of maximum size `chunk_size` tokens and return a list of documents.
-        """
-        text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
-            AutoTokenizer.from_pretrained(tokenizer_name),
-            chunk_size=chunk_size,
-            chunk_overlap=int(chunk_size / 10),
-            add_start_index=True,
-            strip_whitespace=True,
-            separators=MARKDOWN_SEPARATORS,
-        )
-    
-        docs_processed = []
-        for doc in knowledge_base:
-            docs_processed += text_splitter.split_documents([doc])
-    
-        # Remove duplicates
-        unique_texts = {}
-        docs_processed_unique = []
-        for doc in docs_processed:
-            if doc.page_content not in unique_texts:
-                unique_texts[doc.page_content] = True
-                docs_processed_unique.append(doc)
-    
-        return docs_processed_unique
-    
-    def __init__(self, path=""):
 
+    def __init__(self, path=""):
         # Config LangChain
         os.environ["LANGCHAIN_TRACING_V2"] = "true"
         os.environ["LANGCHAIN_API_KEY"] = "ls__9834e6b2ff094d43a28418c9ecea2fd5"
-
-        EMBEDDING_MODEL_NAME = "mixedbread-ai/mxbai-embed-large-v1"
         
         # Load Vector db
         urls = [
@@ -77,34 +61,27 @@ class EndpointHandler():
         
         loader = WebBaseLoader(urls)
         docs = loader.load()
-
-        MARKDOWN_SEPARATORS = [
-            "\n#{1,6} ",
-            "```\n",
-            "\n\\*\\*\\*+\n",
-            "\n---+\n",
-            "\n___+\n",
-            "\n\n",
-            "\n",
-            " ",
-            "",
-        ]
         
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=1000,  # the maximum number of characters in a chunk: we selected this value arbitrarily
-            chunk_overlap=100,  # the number of characters to overlap between chunks
-            add_start_index=True,  # If `True`, includes chunk's start index in metadata
-            strip_whitespace=True,  # If `True`, strips whitespace from the start and end of every document
+        text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
+            AutoTokenizer.from_pretrained(EMBEDDING_MODEL_NAME),
+            chunk_size=512,
+            chunk_overlap=int(512 / 10),
+            add_start_index=True,
+            strip_whitespace=True,
             separators=MARKDOWN_SEPARATORS,
         )
-        
-        docs_processed = text_splitter.split_documents(docs)
 
-        docs_processed = split_documents(
-            512,  # We choose a chunk size adapted to our model
-            docs,
-            tokenizer_name=EMBEDDING_MODEL_NAME,
-        )
+        docs_processed = []
+        for doc in docs:
+            docs_processed += text_splitter.split_documents([doc])
+
+        # Remove duplicates
+        unique_texts = {}
+        docs_processed_unique = []
+        for doc in docs_processed:
+            if doc.page_content not in unique_texts:
+                unique_texts[doc.page_content] = True
+                docs_processed_unique.append(doc)
 
         embedding_model = HuggingFaceEmbeddings(
             model_name=EMBEDDING_MODEL_NAME,
@@ -114,7 +91,7 @@ class EndpointHandler():
         )
         
         self.vectorstore = FAISS.from_documents(
-            docs_processed, embedding_model, distance_strategy=DistanceStrategy.COSINE
+            docs_processed_unique, embedding_model, distance_strategy=DistanceStrategy.COSINE
         )
 
         # Create LLM
@@ -132,7 +109,7 @@ class EndpointHandler():
         # Testing
         # tokenizer.pad_token = tokenizer.eos_token
         
-        READER_LLM = pipeline(
+        self.READER_LLM = pipeline(
             model=model,
             tokenizer=tokenizer,
             task="text-generation",
@@ -180,7 +157,7 @@ class EndpointHandler():
         )
         
         # Redact an answer
-        answer = READER_LLM(final_prompt)[0]["generated_text"]
+        answer = self.READER_LLM(final_prompt)[0]["generated_text"]
 
         return answer