Update app.py

app.py

@@ -4,7 +4,6 @@ from langchain.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
 from langchain.llms import HuggingFaceHub
-from langchain.prompts import PromptTemplate
 from langchain.chains import ConversationalRetrievalChain
 
 # Load the HuggingFace language model and embeddings
@@ -13,19 +12,24 @@ client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 # Initialize the embeddings model for document retrieval
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
-# Set up FAISS as the vector store for document retrieval
-# (Replace 'documents' with your actual document list or corpus)
-texts = ["Document 1", "Document 2", "Document 3"]
-vector_store = FAISS.from_texts(texts, embeddings)
+# Initialize vector_store and retriever as None initially
+vector_store = None
+retriever = None
 
-# Initialize the LangChain RAG system
-retriever = vector_store.as_retriever()
+def update_documents(text_input):
+    global vector_store, retriever
+    # Split the input text into individual documents based on newlines or other delimiters
+    documents = text_input.split("\n")
+    
+    # Update the FAISS vector store with new documents
+    vector_store = FAISS.from_texts(documents, embeddings)
+    
+    # Set the retriever to use the new vector store
+    retriever = vector_store.as_retriever()
+    return f"{len(documents)} documents successfully added to the vector store."
 
-# Set up ConversationalRetrievalChain using LangChain's tools
-rag_chain = ConversationalRetrievalChain.from_llm(
-    HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta"),
-    retriever=retriever
-)
+# Set up ConversationalRetrievalChain
+rag_chain = None
 
 def respond(
     message,
@@ -35,6 +39,18 @@ def respond(
     temperature,
     top_p,
 ):
+    global rag_chain, retriever
+    
+    if retriever is None:
+        return "Please upload or enter documents before asking a question."
+
+    # Create the chain if it hasn't been initialized
+    if rag_chain is None:
+        rag_chain = ConversationalRetrievalChain.from_llm(
+            HuggingFaceHub(repo_id="HuggingFaceH4/zephyr-7b-beta"),
+            retriever=retriever
+        )
+
     # Combine history with the user message
     conversation_history = [{"role": "system", "content": system_message}]
     
@@ -46,22 +62,38 @@ def respond(
 
     conversation_history.append({"role": "user", "content": message})
 
-    # Retrieve documents using the retriever
+    # Retrieve documents and generate response
     response = rag_chain({"question": message, "chat_history": history})
 
-    # Format and return the response
+    # Return the model's response
     return response['answer']
 
+
 # Gradio interface setup
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a helpful assistant.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
-    ],
-)
+demo = gr.Blocks()
+
+with demo:
+    with gr.Row():
+        # Input box for user to add documents
+        doc_input = gr.Textbox(
+            lines=10, placeholder="Enter your documents here, one per line.", label="Input Documents"
+        )
+        upload_button = gr.Button("Upload Documents")
+    
+    with gr.Row():
+        # Chat interface for the RAG system
+        chat = gr.ChatInterface(
+            respond,
+            additional_inputs=[
+                gr.Textbox(value="You are a helpful assistant.", label="System message"),
+                gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+                gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
+                gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+            ],
+        )
+
+    # Bind button to update the document vector store
+    upload_button.click(update_documents, inputs=[doc_input], outputs=gr.Textbox(label="Status"))
 
 if __name__ == "__main__":
     demo.launch()