Update vector_store_retriever.py

vector_store_retriever.py

@@ -1,31 +1,40 @@
 import gradio as gr
-from langchain.document_loaders import TextLoader
+from langchain.document_loaders import DirectoryLoader, PyPDFLoader
 from langchain.vectorstores import Chroma
 from langchain.chains import RetrievalQA
-from langchain.embeddings import HuggingFaceInstructEmbeddings 
+from langchain.embeddings import HuggingFaceInstructEmbeddings
 from langchain.agents import Tool
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.llms import HuggingFacePipeline
+from transformers import LlamaTokenizer, LlamaForCausalLM, pipeline
 
-# Initialize the HuggingFaceInstructEmbeddings
-hf = HuggingFaceInstructEmbeddings(
-  model_name="hkunlp/instructor-large",
-  embed_instruction="Represent the document for retrieval: ",
-  query_instruction="Represent the query for retrieval: "
-)
+# Load and process the text files
+loader = DirectoryLoader('./new_papers/new_papers/', glob="./*.pdf", loader_cls=PyPDFLoader)
+documents = loader.load()
 
-# Example texts for the vector store
-texts=["The meaning of life is to love","The meaning of vacation is to relax","Roses are red.","Hack the planet!"]
+# Splitting the text into chunks
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+texts = text_splitter.split_documents(documents)
 
-# Create a Chroma vector store from the example texts
-db = Chroma.from_texts(texts, hf, collection_name="my-collection")
+# HF Instructor Embeddings
+instructor_embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-xl", model_kwargs={"device": "cuda"})
 
-# Create a RetrievalQA chain
-llm = LLM.from_model("lgaalves/gpt2-dolly")  # Replace with the appropriate LLM model
-docsearcher = RetrievalQA.from_chain_type(
-  llm=llm, 
-  chain_type="stuff",  # Replace with the appropriate chain type
-  return_source_documents=False,
-  retriever=db.as_retriever(search_type="similarity", search_kwargs={"k": 1})
-)
+# Embed and store the texts
+persist_directory = 'db'
+embedding = instructor_embeddings
+vectordb = Chroma.from_documents(documents=texts, embedding=embedding, persist_directory=persist_directory)
+
+# Make a retriever
+retriever = vectordb.as_retriever(search_kwargs={"k": 3})
+
+# Setup LLM for text generation
+tokenizer = LlamaTokenizer.from_pretrained("TheBloke/wizardLM-7B-HF")
+model = LlamaForCausalLM.from_pretrained("TheBloke/wizardLM-7B-HF", load_in_8bit=True, device_map='auto', torch_dtype=torch.float16, low_cpu_mem_usage=True)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=1024, temperature=0, top_p=0.95, repetition_penalty=1.15)
+local_llm = HuggingFacePipeline(pipeline=pipe)
+
+# Make a chain
+qa_chain = RetrievalQA.from_chain_type(llm=local_llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
 
 class VectorStoreRetrieverTool(Tool):
     name = "vectorstore_retriever"
@@ -36,8 +45,8 @@ class VectorStoreRetrieverTool(Tool):
 
     def __call__(self, query: str):
         # Run the query through the RetrievalQA chain
-        response = docsearcher.run(query)
-        return response
+        llm_response = qa_chain(query)
+        return llm_response['result']
 
 # Create the Gradio interface using the HuggingFaceTool
 tool = gr.Interface(