Update app.py

app.py

@@ -1,30 +1,20 @@
 import os
 import gradio as gr
 from huggingface_hub import InferenceClient
+from langchain_community.vectorstores import Chroma
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain_community.document_loaders import PyPDFLoader, UnstructuredFileLoader, CSVLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceBgeEmbeddings
-from langchain.document_loaders import PyPDFLoader, UnstructuredFileLoader, CSVLoader
-from langchain.chains import RetrievalQA
-from langchain.prompts import PromptTemplate
 
-# # Initialize the Zephyr client
-# client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-from huggingface_hub import InferenceClient
-
-# Access the Hugging Face token from environment variables
+# Load Hugging Face API token
 HF_API_TOKEN = os.getenv("HF_API_TOKEN")
-
 if not HF_API_TOKEN:
     raise ValueError("Hugging Face API token is not set in environment variables.")
 
-# Initialize the client with the token
+# Initialize Zephyr client
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta", token=HF_API_TOKEN)
 
-
-
-# Function to load documents based on file type
+# Load documents based on file type
 def load_documents(file_path):
     if file_path.endswith(".pdf"):
         loader = PyPDFLoader(file_path)
@@ -34,54 +24,50 @@ def load_documents(file_path):
         loader = CSVLoader(file_path)
     else:
         raise ValueError("Unsupported file format")
-    
     documents = loader.load()
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
     return text_splitter.split_documents(documents)
 
-# Function to create or update vector store
+# Create vector store
 def create_vector_store(documents, persist_dir="vector_db"):
     embeddings = HuggingFaceBgeEmbeddings(
-        model_name="BAAI/bge-large-en", 
-        model_kwargs={"device": "cpu"}
+        model_name="BAAI/bge-large-en",
+        model_kwargs={"device": "cpu"},
     )
     vector_store = Chroma.from_documents(documents, embeddings, persist_directory=persist_dir)
     return vector_store
 
-# Function to handle user queries
-def respond(message, history, system_message, max_tokens, temperature, top_p, retriever):
+# Initialize retriever and vector store
+persist_dir = "vector_db"
+retriever = None  # Will be dynamically updated
+
+# Handle queries and uploads
+def handle_query(message, history, system_message, max_tokens, temperature, top_p, file=None):
+    global retriever
+    if file:  # If a file is uploaded, process it
+        documents = load_documents(file.name)
+        vector_store = create_vector_store(documents, persist_dir)
+        retriever = vector_store.as_retriever()
+    if not retriever:
+        return "No documents uploaded yet. Please upload a file first."
+    
     # Retrieve relevant context
     relevant_docs = retriever.get_relevant_documents(message)
     context = "\n".join([doc.page_content for doc in relevant_docs])
-
-    # Format the prompt
-    prompt_template = """
-    Use the following context to answer the user's question. 
-    If you don't know the answer, say "I don't know."
-
+    
+    # Build the prompt
+    prompt = f"""
+    Use the following context to answer the user's question.
     Context:
     {context}
-
+    
     Question:
-    {question}
-
-    Answer:
-    """
-    formatted_prompt = prompt_template.format(context=context, question=message)
-
-    # Build conversational history
-    messages = [{"role": "system", "content": system_message}]
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-    messages.append({"role": "user", "content": formatted_prompt})
-
-    # Stream response from Zephyr
+    {message}
+    Answer:"""
+    
     response = ""
     for msg in client.chat_completion(
-        messages=messages,
+        messages=[{"role": "system", "content": system_message}, {"role": "user", "content": prompt}],
         max_tokens=max_tokens,
         stream=True,
         temperature=temperature,
@@ -91,37 +77,9 @@ def respond(message, history, system_message, max_tokens, temperature, top_p, re
         response += token
         yield response
 
-# Initialize the vector store
-persist_dir = "vector_db"
-retriever = None  # Will be initialized dynamically
-
-def handle_query(message, history, system_message, max_tokens, temperature, top_p, file=None):
-    global retriever
-    if file:  # Process uploaded file
-        documents = load_documents(file.name)
-        vector_store = create_vector_store(documents, persist_dir)
-        retriever = vector_store.as_retriever()
-    if not retriever:
-        return "No documents have been uploaded yet. Please upload a file to provide context."
-    return respond(message, history, system_message, max_tokens, temperature, top_p, retriever)
-
 # Gradio app setup
-# demo = gr.ChatInterface(
-#     fn=handle_query,
-#     additional_inputs=[
-#         gr.File(label="Upload File", type="file"),
-#         gr.Textbox(value="You are a knowledgeable assistant.", label="System Message"),
-#         gr.Slider(1, 2048, step=1, value=512, label="Max Tokens"),
-#         gr.Slider(0.1, 4.0, step=0.1, value=0.7, label="Temperature"),
-#         gr.Slider(0.1, 1.0, step=0.05, value=0.95, label="Top-p"),
-#     ],
-#     outputs="text",
-#     title="RAG with Zephyr-7B",
-#     description="A Retrieval-Augmented Generation chatbot powered by Zephyr-7B and Chroma vector database.",
-# )
-
 demo = gr.Interface(
-    fn=handle_uploaded_file,  # Handle uploaded files
+    fn=handle_query,
     inputs=[
         gr.File(label="Upload Document"),
         gr.Textbox(value="You are a knowledgeable assistant.", label="System Message"),
@@ -135,7 +93,5 @@ demo = gr.Interface(
     description="Upload documents and ask questions using RAG.",
 )
 
-
-
 if __name__ == "__main__":
     demo.launch()