Update app.py

app.py

@@ -1,30 +1,36 @@
 import os
-import fitz
+import fitz  # PyMuPDF
 from docx import Document
 from sentence_transformers import SentenceTransformer
-from langchain_community.llms import HuggingFaceEndpoint  # Might need update (optional)
 from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from nltk.tokenize import sent_tokenize  # Import for sentence segmentation
+from nltk.tokenize import sent_tokenize
+import torch
+import gradio as gr
 
-# Function to extract text from a PDF file (same as before)
+# Function to extract text from a PDF file
 def extract_text_from_pdf(pdf_path):
-    # ... (implementation)
+    text = ""
+    doc = fitz.open(pdf_path)
+    for page in doc:
+        text += page.get_text()
+    return text
 
-# Function to extract text from a Word document (same as before)
+# Function to extract text from a Word document
 def extract_text_from_docx(docx_path):
-    # ... (implementation)
+    doc = Document(docx_path)
+    text = "\n".join([para.text for para in doc.paragraphs])
+    return text
 
-# Initialize the embedding model (same as before)
+# Initialize the embedding model
 embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 
-# Hugging Face API token (same as before)
+# Hugging Face API token
 api_token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
 if not api_token:
     raise ValueError("HUGGINGFACEHUB_API_TOKEN environment variable is not set")
 
-# Define RAG models (same as before)
+# Define RAG models
 generator_model_name = "facebook/bart-base"
 retriever_model_name = "facebook/bart-base"  # Can be the same as generator
 generator = AutoModelForSeq2SeqLM.from_pretrained(generator_model_name)
@@ -32,7 +38,7 @@ generator_tokenizer = AutoTokenizer.from_pretrained(generator_model_name)
 retriever = AutoModelForSeq2SeqLM.from_pretrained(retriever_model_name)
 retriever_tokenizer = AutoTokenizer.from_pretrained(retriever_model_name)
 
-# Load or create FAISS index (using LangChain)
+# Load or create FAISS index
 index_path = "faiss_index.pkl"
 if os.path.exists(index_path):
     with open(index_path, "rb") as f:
@@ -49,7 +55,7 @@ def preprocess_text(text):
     sentences = sent_tokenize(text)
     return sentences
 
-def upload_files(state, files):
+def upload_files(files):
     global index
     try:
         for file_path in files:
@@ -60,7 +66,7 @@ def upload_files(state, files):
             else:
                 return {"error": "Unsupported file format"}
 
-            # Preprocess text (call the new function)
+            # Preprocess text
             sentences = preprocess_text(text)
 
             # Encode sentences and add to FAISS index
@@ -70,36 +76,20 @@ def upload_files(state, files):
         return {"message": "Files processed successfully"}
     except Exception as e:
         print(f"Error processing files: {e}")
-        return {"error": "Error processing files"}  # Provide informative error message
+        return {"error": "Error processing files"}
 
 def process_and_query(state, files, question):
-    # State management for conversation history (similar to previous example)
-    # ...
-
-    # Handle file upload (using upload_files function)
-    if files:
-        upload_result = upload_files(state, files)
-        if "error" in upload_result:
-            return upload_result  # Return error message from upload_files if any
-
-    # Handle user question and generate response using RAG models if question and state.
-    def process_and_query(state, files, question):
-    # State management for conversation history (similar to previous example)
-    # ...
-
-    # Handle file upload (using upload_files function)
     if files:
-        upload_result = upload_files(state, files)
+        upload_result = upload_files(files)
         if "error" in upload_result:
-            return upload_result  # Return error message from upload_files if any
+            return upload_result
 
-    # Handle user question and generate response using RAG models
-    if question and state["processed_text"]:
+    if question:
         # Preprocess the question
         question_embedding = embedding_model.encode([question])
 
         # Use retriever model to retrieve relevant passages
-        with torch.no_grad():  # Disable gradient calculation for efficiency
+        with torch.no_grad():
             retriever_outputs = retriever(**retriever_tokenizer(question, return_tensors="pt"))
             retriever_hidden_states = retriever_outputs.hidden_states[-1]  # Last hidden state
 
@@ -110,7 +100,6 @@ def process_and_query(state, files, question):
         retrieved_passages = [state["processed_text"].split("\n")[i] for i in retrieved_ids.flatten()]
 
         # Use generator model to generate response based on question and retrieved passages
-        # Combine question embedding with retrieved passages (consider weighting or attention mechanism)
         combined_input = torch.cat([question_embedding, embedding_model.encode(retrieved_passages)], dim=0)
         with torch.no_grad():
             generator_outputs = generator(**generator_tokenizer(combined_input, return_tensors="pt"))
@@ -119,23 +108,26 @@ def process_and_query(state, files, question):
         # Update conversation history
         state["conversation"].append({"question": question, "answer": generated_text})
 
-        return state  # Return the updated state with conversation history
+        return {"message": generated_text, "conversation": state["conversation"]}
+
+    return {"error": "No question provided"}
 
 # Create Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("## Document Upload and Query System")
-    
+
     with gr.Tab("Upload Files"):
         upload = gr.File(file_count="multiple", label="Upload PDF or DOCX files")
         upload_button = gr.Button("Upload")
         upload_output = gr.Textbox()
         upload_button.click(fn=upload_files, inputs=upload, outputs=upload_output)
-    
+
     with gr.Tab("Query"):
         query = gr.Textbox(label="Enter your query")
         query_button = gr.Button("Search")
         query_output = gr.Textbox()
-        query_button.click(fn=query_text, inputs=query, outputs=query_output)
+        query_button.click(fn=process_and_query, inputs=[query], outputs=query_output)
 
 demo.launch()
 
+