Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import sqlite3
 import requests
-import PyPDF2
+import fitz  # PyMuPDF
 import faiss
 import numpy as np
 from sentence_transformers import SentenceTransformer
@@ -19,11 +19,11 @@ def query_huggingface(payload):
 
 # Function to extract text from PDF
 def extract_text_from_pdf(pdf_file):
-    pdf_reader = PyPDF2.PdfReader(pdf_file)
     text = ""
-    for page_num in range(len(pdf_reader.pages)):
-        page = pdf_reader.pages[page_num]
-        text += page.extract_text()
+    pdf_document = fitz.open(stream=pdf_file.read(), filetype="pdf")
+    for page_num in range(len(pdf_document)):
+        page = pdf_document.load_page(page_num)
+        text += page.get_text()
     return text
 
 # Initialize SQLite database
@@ -60,6 +60,9 @@ def get_context():
 # Function to create or update the FAISS index
 def update_faiss_index():
     contexts = get_context()
+    if len(contexts) == 0:
+        return None, contexts
+
     embeddings = model.encode(contexts, convert_to_tensor=True)
     index = faiss.IndexFlatL2(embeddings.shape[1])
     index.add(embeddings.cpu().numpy())
@@ -67,6 +70,9 @@ def update_faiss_index():
 
 # Retrieve relevant context from the FAISS index
 def retrieve_relevant_context(index, contexts, query, top_k=5):
+    if index is None or len(contexts) == 0:
+        return []
+    
     query_embedding = model.encode([query], convert_to_tensor=True).cpu().numpy()
     distances, indices = index.search(query_embedding, top_k)
     relevant_contexts = [contexts[i] for i in indices[0]]
@@ -77,49 +83,25 @@ init_db()
 model = SentenceTransformer('all-MiniLM-L6-v2')
 faiss_index, context_list = update_faiss_index()
 
-# Function to handle chatbot responses
-def chatbot_response(question):
+# Gradio interface
+def chatbot(question):
     relevant_contexts = retrieve_relevant_context(faiss_index, context_list, question)
     user_input = f"question: {question} context: {' '.join(relevant_contexts)}"
     response = query_huggingface({"inputs": user_input})
     response_text = response.get("generated_text", "Sorry, I couldn't generate a response.")
     return response_text
 
-# Function to handle PDF uploads
-def handle_pdf_upload(pdf_file):
-    context = extract_text_from_pdf(pdf_file)
-    add_context(pdf_file.name, context)
-    faiss_index, context_list = update_faiss_index()  # Update FAISS index
-    return f"Context from {pdf_file.name} added to the database."
-
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("# Storage Warehouse Customer Service Chatbot")
-
-    with gr.Row():
-        with gr.Column(scale=4):
-            with gr.Box():
-                pdf_upload = gr.File(label="Upload PDF", file_types=["pdf"], interactive=True)
-                upload_button = gr.Button("Upload")
-                upload_status = gr.Textbox(label="Upload Status")
-
-            def handle_upload(files):
-                for file in files:
-                    result = handle_pdf_upload(file.name)
-                    upload_status.value = result
-
-            upload_button.click(fn=handle_upload, inputs=pdf_upload, outputs=upload_status)
-
-        with gr.Column(scale=8):
-            chatbot = gr.Chatbot(label="Chatbot")
-            question = gr.Textbox(label="Your question here:")
-            submit_button = gr.Button("Submit")
-
-            def handle_chat(user_input):
-                bot_response = chatbot_response(user_input)
-                return gr.Chatbot.update([[user_input, bot_response]])
+# File upload function
+def upload_pdf(file):
+    context = extract_text_from_pdf(file)
+    add_context(file.name, context)
+    global faiss_index, context_list
+    faiss_index, context_list = update_faiss_index()
+    return "PDF content added to context."
 
-            submit_button.click(fn=handle_chat, inputs=question, outputs=chatbot)
+# Gradio interface
+iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", title="Storage Warehouse Customer Service Chatbot")
+file_upload = gr.Interface(fn=upload_pdf, inputs="file", outputs="text", title="Upload PDF for Context")
 
-if __name__ == "__main__":
-    demo.launch()
+app = gr.TabbedInterface([iface, file_upload], ["Chatbot", "Upload PDF"])
+app.launch()