Update app.py

app.py

@@ -1,22 +1,18 @@
 import os
-import faiss
-import numpy as np
-from docx import Document
+import gradio as gr
+import fitz  # PyMuPDF for PDF text extraction
+from docx import Document  # python-docx for DOCX text extraction
 from sentence_transformers import SentenceTransformer
 from langchain_community.vectorstores import FAISS
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from nltk.tokenize import sent_tokenize
 import torch
-import gradio as gr
 import pickle
 import nltk
 
-# Download NLTK punkt resource if not already downloaded
-try:
-    nltk.data.find('tokenizers/punkt')
-except LookupError:
-    nltk.download('punkt')
+# Download NLTK punkt tokenizer data if not already downloaded
+nltk.download('punkt', quiet=True)
 
 # Function to extract text from a PDF file
 def extract_text_from_pdf(pdf_path):
@@ -40,63 +36,46 @@ def extract_text_from_docx(docx_path):
         print(f"Error extracting text from DOCX: {e}")
     return text
 
-
-# Initialize the embedding model
+# Initialize the SentenceTransformer model for embeddings
 embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 
-# 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
-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)
-generator_tokenizer = AutoTokenizer.from_pretrained(generator_model_name)
-retriever = AutoModelForSeq2SeqLM.from_pretrained(retriever_model_name)
-retriever_tokenizer = AutoTokenizer.from_pretrained(retriever_model_name)
-
-# Initialize FAISS index using LangChain
+# Initialize the HuggingFaceEmbeddings for LangChain
 hf_embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
 
-# Load or create FAISS index
+# Initialize the FAISS index
 index_path = "faiss_index.pkl"
 if os.path.exists(index_path):
     with open(index_path, "rb") as f:
         faiss_index = pickle.load(f)
         print("Loaded FAISS index from faiss_index.pkl")
 else:
-    # Initialize a new FAISS index, e.g., IndexIVFFlat
-    d = 384  # Embedding dimension
-    nlist = 100  # Number of clusters
-    quantizer = faiss.IndexFlatL2(d)
-    faiss_index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
-    faiss_index.train(np.array([]))  # Optional: Train index if needed
-    faiss_index.add(np.array([]))  # Initialize index
+    # Initialize FAISS index using LangChain
+    faiss_index = FAISS(embedding_function=hf_embeddings)
 
 def preprocess_text(text):
     sentences = sent_tokenize(text)
     return sentences
 
 def upload_files(files):
-    global faiss_index
     try:
         for file in files:
-            if file.name.endswith('.pdf'):
-                text = extract_text_from_pdf(file.name)
-            elif file.name.endswith('.docx'):
-                text = extract_text_from_docx(file.name)
+            if isinstance(file, str):  # Assuming `file` is a string (file path)
+                if file.endswith('.pdf'):
+                    text = extract_text_from_pdf(file)
+                elif file.endswith('.docx'):
+                    text = extract_text_from_docx(file)
+                else:
+                    return {"error": "Unsupported file format"}
             else:
-                return {"error": "Unsupported file format"}
+                return {"error": "Invalid file format: expected a string"}
 
             # Preprocess text
             sentences = preprocess_text(text)
 
             # Encode sentences and add to FAISS index
             embeddings = embedding_model.encode(sentences)
-            for embedding in embeddings:
-                faiss_index.add(np.expand_dims(embedding, axis=0))
+            for sentence, embedding in zip(sentences, embeddings):
+                faiss_index.add_sentence(sentence, embedding)
 
         # Save the updated index
         with open(index_path, "wb") as f:
@@ -113,15 +92,18 @@ def process_and_query(state, files, question):
         if "error" in upload_result:
             return upload_result
 
-    if question and question.strip():  # Check if question is not empty
+    if question:
         # Preprocess the question
         question_embedding = embedding_model.encode([question])
 
         # Search the FAISS index for similar passages
-        D, I = faiss_index.search(np.array([question_embedding]), 5)  # Retrieve top 5 passages
-        retrieved_passages = []
-        for i in I[0]:
-            retrieved_passages.append(faiss_index.reconstruct(i).decode('utf-8'))
+        retrieved_results = faiss_index.similarity_search(question, k=5)  # Retrieve top 5 passages
+        retrieved_passages = [result['text'] for result in retrieved_results]
+
+        # Initialize RAG generator model
+        generator_model_name = "facebook/bart-base"
+        generator = AutoModelForSeq2SeqLM.from_pretrained(generator_model_name)
+        generator_tokenizer = AutoTokenizer.from_pretrained(generator_model_name)
 
         # Use generator model to generate response based on question and retrieved passages
         combined_input = question + " ".join(retrieved_passages)
@@ -151,6 +133,7 @@ with gr.Blocks() as demo:
         query = gr.Textbox(label="Enter your query")
         query_button = gr.Button("Search")
         query_output = gr.Textbox()
-        query_button.click(fn=process_and_query, inputs=[upload_output, query], outputs=query_output)
+        query_button.click(fn=process_and_query, inputs=[query], outputs=query_output)
+
+demo.launch()
 
-demo.launch()