Update app.py

app.py

@@ -2,11 +2,13 @@ import os
 import fitz  # PyMuPDF
 from docx import Document
 from sentence_transformers import SentenceTransformer
-from langchain_community.vectorstores import FAISS
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from nltk.tokenize import sent_tokenize
 import torch
 import gradio as gr
+import pickle
 
 # Function to extract text from a PDF file
 def extract_text_from_pdf(pdf_path):
@@ -38,31 +40,31 @@ 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
+embedding_dimension = embedding_model.get_sentence_embedding_dimension()
+faiss_index = FAISS(HuggingFaceEmbeddings(embedding_model), dimension=embedding_dimension)
+
 # Load or create FAISS index
 index_path = "faiss_index.pkl"
 if os.path.exists(index_path):
     with open(index_path, "rb") as f:
-        index = FAISS.load(f)
+        faiss_index = pickle.load(f)
         print("Loaded FAISS index from faiss_index.pkl")
 else:
-    # Create a new FAISS index if it doesn't exist
-    index = FAISS(embedding_dimension=embedding_model.get_sentence_embedding_dimension())
-    with open(index_path, "wb") as f:
-        FAISS.save(index, f)
-        print("Created new FAISS index and saved to faiss_index.pkl")
+    print("Created new FAISS index")
 
 def preprocess_text(text):
     sentences = sent_tokenize(text)
     return sentences
 
 def upload_files(files):
-    global index
+    global faiss_index
     try:
-        for file_path in files:
-            if file_path.endswith('.pdf'):
-                text = extract_text_from_pdf(file_path)
-            elif file_path.endswith('.docx'):
-                text = extract_text_from_docx(file_path)
+        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)
             else:
                 return {"error": "Unsupported file format"}
 
@@ -71,7 +73,11 @@ def upload_files(files):
 
             # Encode sentences and add to FAISS index
             embeddings = embedding_model.encode(sentences)
-            index.add(embeddings)
+            faiss_index.add_texts(sentences, embeddings)
+
+        # Save the updated index
+        with open(index_path, "wb") as f:
+            pickle.dump(faiss_index, f)
 
         return {"message": "Files processed successfully"}
     except Exception as e:
@@ -88,22 +94,18 @@ def process_and_query(state, files, question):
         # Preprocess the question
         question_embedding = embedding_model.encode([question])
 
-        # Use retriever model to retrieve relevant passages
-        with torch.no_grad():
-            retriever_outputs = retriever(**retriever_tokenizer(question, return_tensors="pt"))
-            retriever_hidden_states = retriever_outputs.hidden_states[-1]  # Last hidden state
-
-        # Search the FAISS index for similar passages based on retrieved hidden states
-        distances, retrieved_ids = index.search(retriever_hidden_states.cpu().numpy(), k=5)  # Retrieve top 5 passages
+        # Search the FAISS index for similar passages
+        distances, retrieved_ids = faiss_index.similarity_search_with_score(question_embedding, k=5)  # Retrieve top 5 passages
 
         # Get the retrieved passages from the document text
         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
-        combined_input = torch.cat([question_embedding, embedding_model.encode(retrieved_passages)], dim=0)
+        combined_input = question + " ".join(retrieved_passages)
+        inputs = generator_tokenizer(combined_input, return_tensors="pt")
         with torch.no_grad():
-            generator_outputs = generator(**generator_tokenizer(combined_input, return_tensors="pt"))
-            generated_text = generator_tokenizer.decode(generator_outputs.sequences.squeeze())
+            generator_outputs = generator.generate(**inputs)
+            generated_text = generator_tokenizer.decode(generator_outputs[0], skip_special_tokens=True)
 
         # Update conversation history
         state["conversation"].append({"question": question, "answer": generated_text})
@@ -131,3 +133,4 @@ with gr.Blocks() as demo:
 demo.launch()
 
 
+