app.py

import os
import faiss
import numpy as np
import PyPDF2
import io
from docx import Document
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

nltk.download('punkt')

# Function to extract text from a PDF file
def extract_text_from_pdf(pdf_data):
    text = ""
    try:
        pdf_reader = PyPDF2.PdfReader(io.BytesIO(pdf_data))
        for page in pdf_reader.pages:
            text += page.extract_text()
    except Exception as e:
        print(f"Error extracting text from PDF: {e}")
    return text

# Function to extract text from a Word document
def extract_text_from_docx(docx_data):
    text = ""
    try:
        doc = Document(io.BytesIO(docx_data))
        text = "\n".join([para.text for para in doc.paragraphs])
    except Exception as e:
        print(f"Error extracting text from DOCX: {e}")
    return text

# Initialize the embedding model
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
hf_embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')

# Load or create 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:
    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.data)
            elif file.name.endswith('.docx'):
                text = extract_text_from_docx(file.data)
            else:
                return {"error": "Unsupported file format"}

            # 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))

        # 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:
        print(f"Error processing files: {e}")
        return {"error": str(e)}  # Provide informative error message

def process_and_query(state, question):
    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), k=5)
        retrieved_passages = [faiss_index.index_to_text(i) for i in I[0]]

        # Use generator model to generate response based on question and retrieved passages
        prompt_template = """
        Answer the question as detailed as possible from the provided context,
        make sure to provide all the details, if the answer is not in
        provided context just say, "answer is not available in the context",
        don't provide the wrong answer

        Context:\n{context}\n
        Question:\n{question}\n
        Answer:
        """
        combined_input = prompt_template.format(context=' '.join(retrieved_passages), question=question)
        inputs = generator_tokenizer(combined_input, return_tensors="pt")
        with torch.no_grad():
            generator_outputs = generator.generate(**inputs)
            generated_text = generator_tokenizer.decode(generator_outputs[0], skip_special_tokens=True)

        # Update conversation history
        state.append({"question": question, "answer": generated_text})

        return {"message": generated_text, "conversation": state}

    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()
        state = gr.State([])
        query_button.click(fn=process_and_query, inputs=[state, query], outputs=query_output)

demo.launch()