Update app.py

app.py

@@ -1,27 +1,91 @@
 import gradio as gr
 import torch
+import faiss
+import numpy as np
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+import pdfplumber
+import docx
 
-# Load RAG model
+# Load RAG Model
 model_name = "facebook/rag-sequence-nq"
 tokenizer = RagTokenizer.from_pretrained(model_name)
 retriever = RagRetriever.from_pretrained(model_name, index_name="exact", use_dummy_dataset=True)
 model = RagSequenceForGeneration.from_pretrained(model_name, retriever=retriever)
 
-# Function to process uploaded document
-def process_file(file):
+# FAISS Vector Store
+dimension = 768  # Default embedding size for transformers
+index = faiss.IndexFlatL2(dimension)  # L2 distance-based index
+stored_docs = []  # To store document texts alongside vectors
+
+# Function to extract text from uploaded files
+def extract_text(file):
     if file is None:
         return "Please upload a document."
     
-    file_text = file.decode("utf-8")
-    return file_text
+    file_name = file.name
+    file_ext = file_name.split('.')[-1].lower()
+    text = ""
+
+    if file_ext == "txt":
+        text = file.read().decode("utf-8")
+    
+    elif file_ext == "pdf":
+        with pdfplumber.open(file) as pdf:
+            for page in pdf.pages:
+                text += page.extract_text() + "\n"
+
+    elif file_ext == "docx":
+        doc = docx.Document(file)
+        for para in doc.paragraphs:
+            text += para.text + "\n"
+
+    else:
+        return "Unsupported file format! Please upload TXT, PDF, or DOCX."
+
+    # Store document in FAISS index
+    store_in_faiss(text.strip())
+    
+    return text.strip()
 
-# Function to answer questions using RAG
+# Function to store document in FAISS
+def store_in_faiss(document):
+    global index, stored_docs
+    if not document.strip():
+        return
+
+    # Tokenize and get embeddings
+    inputs = tokenizer(document, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        embeddings = model.rag.retriever(input_ids=inputs["input_ids"]).cpu().numpy()
+    
+    # Add embeddings to FAISS
+    index.add(embeddings)
+    stored_docs.append(document)
+
+# Function to retrieve top relevant document from FAISS
+def retrieve_relevant_doc(query):
+    if index.ntotal == 0:
+        return ""
+
+    # Tokenize query and get embeddings
+    inputs = tokenizer(query, return_tensors="pt", truncation=True, max_length=512)
+    with torch.no_grad():
+        query_embedding = model.rag.retriever(input_ids=inputs["input_ids"]).cpu().numpy()
+
+    # Search in FAISS
+    _, top_idx = index.search(query_embedding, k=1)
+    return stored_docs[top_idx[0][0]]
+
+# Function to answer questions using RAG with FAISS
 def answer_question(document, question):
     if not document.strip():
         return "Please provide document content."
 
-    inputs = tokenizer(question, document, return_tensors="pt", truncation=True)
+    # Retrieve best-matching document
+    relevant_doc = retrieve_relevant_doc(question)
+
+    # Use RAG model for answer generation
+    inputs = tokenizer(question, relevant_doc, return_tensors="pt", truncation=True)
     with torch.no_grad():
         generated = model.generate(**inputs)
     answer = tokenizer.batch_decode(generated, skip_special_tokens=True)[0]
@@ -30,20 +94,26 @@ def answer_question(document, question):
 
 # Gradio UI
 with gr.Blocks() as app:
-    gr.Markdown("# 📄 Advanced RAG NLP Document Editor")
+    gr.Markdown("# 📄 Advanced RAG NLP Document Editor with FAISS")
 
     # File Uploader
-    file_input = gr.File(label="Upload Document (TXT only)", type="binary")
-    file_output = gr.Textbox(label="Extracted Text", lines=10)
+    file_input = gr.File(label="Upload Document (TXT, PDF, DOCX)", type="file")
+    file_output = gr.Textbox(label="Extracted Text (Editable)", lines=12)
+
+    file_input.change(extract_text, inputs=file_input, outputs=file_output)
+
+    # Editable Text Editor Canvas
+    editor = gr.Textbox(label="Editor Canvas (Modify Text Before Asking)", lines=12)
 
-    file_input.change(process_file, inputs=file_input, outputs=file_output)
+    # Update editor with extracted text
+    file_output.change(lambda x: x, inputs=file_output, outputs=editor)
 
     # Question Answering
     question_input = gr.Textbox(label="Ask a Question")
     answer_output = gr.Textbox(label="Answer", lines=2)
     
     submit_btn = gr.Button("Get Answer")
-    submit_btn.click(answer_question, inputs=[file_output, question_input], outputs=answer_output)
+    submit_btn.click(answer_question, inputs=[editor, question_input], outputs=answer_output)
 
 # Launch in Hugging Face Spaces
 app.launch()