Update app.py

app.py

@@ -1,44 +1,41 @@
 import gradio as gr
+import os
 from langchain.vectorstores import Chroma
 from langchain.document_loaders import PyPDFLoader
-from langchain.embeddings import HuggingFaceInstructEmbeddings
+from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings
 
-# Initialize the HuggingFaceInstructEmbeddings
-hf = HuggingFaceInstructEmbeddings(
-    model_name="gpt2",
-    embed_instruction="Represent the document for retrieval: ",
-    query_instruction="Represent the query for retrieval: "
+# Use Hugging Face Inference API embeddings
+inference_api_key = os.environ['hf']
+api_hf_embeddings = HuggingFaceInferenceAPIEmbeddings(
+    api_key=inference_api_key,
+    model_name="sentence-transformers/all-MiniLM-l6-v2"
 )
-# Add a padding token to the tokenizer
-hf.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
 
 # Load and process the PDF files
 loader = PyPDFLoader("./new_papers/ReACT.pdf")
-
-#loader = PyPDFLoader('./new_papers/', glob="./*.pdf")
 documents = loader.load()
 
-# Create a Chroma vector store from the PDF documents
-db = Chroma.from_documents(documents, hf, collection_name="my-collection")
+# Create Chroma vector store for API embeddings
+api_db = Chroma.from_documents(documents, api_hf_embeddings, collection_name="api-collection")
 
 class PDFRetrievalTool:
-    def __init__(self):
-        self.retriever = db.as_retriever(search_kwargs={"k": 1})
+    def __init__(self, retriever):
+        self.retriever = retriever
 
     def __call__(self, query):
         # Run the query through the retriever
         response = self.retriever.run(query)
         return response['result']
 
-# Create the Gradio interface using the PDFRetrievalTool
-tool = gr.Interface(
-    PDFRetrievalTool(),
+# Create Gradio interface for the API retriever
+api_tool = gr.Interface(
+    PDFRetrievalTool(api_db.as_retriever(search_kwargs={"k": 1})),
     inputs=gr.Textbox(),
     outputs=gr.Textbox(),
     live=True,
-    title="PDF Retrieval Tool",
-    description="This tool indexes PDF documents and retrieves relevant answers based on a given query.",
+    title="API PDF Retrieval Tool",
+    description="This tool indexes PDF documents and retrieves relevant answers based on a given query (HF Inference API Embeddings).",
 )
 
 # Launch the Gradio interface
-tool.launch()
+api_tool.launch()