Update vector_store_retriever.py

vector_store_retriever.py

@@ -1,59 +1,39 @@
 import gradio as gr
-from langchain.document_loaders import DirectoryLoader, PyPDFLoader
 from langchain.vectorstores import Chroma
-from langchain.chains import RetrievalQA
+from langchain.document_loaders import PyPDFLoader
 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
 
-# Load and process the text files
-loader = DirectoryLoader('./new_papers/', glob="./*.pdf", loader_cls=PyPDFLoader)
-documents = loader.load()
-
-# Splitting the text into chunks
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-texts = text_splitter.split_documents(documents)
-
-# HF Instructor Embeddings
-instructor_embeddings = HuggingFaceInstructEmbeddings(model_name="hkunlp/instructor-xl", model_kwargs={"device": "cuda"})
-
-# 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)
+# Initialize the HuggingFaceInstructEmbeddings
+hf = HuggingFaceInstructEmbeddings(
+    model_name="hkunlp/instructor-large",
+    embed_instruction="Represent the document for retrieval: ",
+    query_instruction="Represent the query for retrieval: "
+)
 
-# Make a chain
-qa_chain = RetrievalQA.from_chain_type(llm=local_llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
+# Load and process the PDF files
+loader = PyPDFLoader('./new_papers/new_papers/', glob="./*.pdf")
+documents = loader.load()
 
-class VectorStoreRetrieverTool(Tool):
-    name = "vectorstore_retriever"
-    description = "This tool uses LangChain's RetrievalQA to find relevant answers from a vector store based on a given query."
+# Create a Chroma vector store from the PDF documents
+db = Chroma.from_documents(documents, hf, collection_name="my-collection")
 
-    inputs = ["text"]
-    outputs = ["text"]
+class VectoreStoreRetrievalTool:
+    def __init__(self):
+        self.retriever = db.as_retriever(search_kwargs={"k": 1})
 
-    def __call__(self, query: str):
-        # Run the query through the RetrievalQA chain
-        llm_response = qa_chain(query)
-        return llm_response['result']
+    def __call__(self, query):
+        # Run the query through the retriever
+        response = self.retriever.run(query)
+        return response['result']
 
-# Create the Gradio interface using the HuggingFaceTool
+# Create the Gradio interface using the PDFRetrievalTool
 tool = gr.Interface(
-    VectorStoreRetrieverTool(),
+    PDFRetrievalTool(),
+    inputs=gr.Textbox(),
+    outputs=gr.Textbox(),
     live=True,
-    title="LangChain-Application: Vectorstore-Retriever",
-    description="This tool uses LangChain's RetrievalQA to find relevant answers from a vector store based on a given query.",
+    title="PDF Retrieval Tool",
+    description="This tool indexes PDF documents and retrieves relevant answers based on a given query.",
 )
 
 # Launch the Gradio interface