Update app.py

app.py

@@ -8,6 +8,7 @@ from nltk.tokenize import sent_tokenize
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 from sentence_transformers import SentenceTransformer
 import gradio as gr
+import torch
 
 # Download NLTK punkt tokenizer if not already downloaded
 import nltk
@@ -16,8 +17,25 @@ nltk.download('punkt')
 # Initialize Sentence Transformer model for embeddings
 embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
 
+# Initialize 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")
+
+# Initialize RAG models from Hugging Face
+generator_model_name = "facebook/bart-base"
+retriever_model_name = "facebook/bart-base"
+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
-faiss_index = None  # Initialize or load your FAISS index as needed
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+
+hf_embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
+faiss_index = FAISS(embedding_function=hf_embeddings)
 
 # Function to extract text from a PDF file
 def extract_text_from_pdf(pdf_data):
@@ -45,7 +63,7 @@ def preprocess_text(text):
     sentences = sent_tokenize(text)
     return sentences
 
-# Function to handle file uploads
+# Function to handle file uploads and update FAISS index
 def upload_files(files):
     global faiss_index
     try:
@@ -82,13 +100,35 @@ def upload_files(files):
         print(f"Error processing files: {e}")
         return {"error": str(e)}  # Provide informative error message
 
-# Function to process queries
+# Function to process queries using RAG model
 def process_and_query(state, question):
     if question:
         try:
-            # Placeholder response based on query processing
-            response_message = "Placeholder response based on query processing"
-            return {"message": response_message, "conversation": state}
+            # Search the FAISS index for similar passages
+            question_embedding = embedding_model.encode([question])
+            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}
         except Exception as e:
             print(f"Error processing query: {e}")
             return {"error": str(e)}
@@ -97,21 +137,29 @@ def process_and_query(state, question):
 
 # Define the Gradio interface
 def main():
+    upload_tab = gr.Interface(
+        fn=upload_files,
+        inputs=gr.inputs.File(label="Upload PDF or DOCX files", multiple=True),
+        outputs=gr.outputs.Text(label="Upload Status", default="No file uploaded yet", type="textbox"),
+        live=True,
+        capture_session=True
+    )
+
+    query_tab = gr.Interface(
+        fn=process_and_query,
+        inputs=gr.inputs.Textbox(label="Enter your query"),
+        outputs=gr.outputs.Textbox(label="Query Response", default="No query processed yet", type="textbox"),
+        live=True,
+        capture_session=True
+    )
+
     gr.Interface(
-        fn=None,  # Replace with your function that handles interface logic
-        inputs=gr.Interface.Layout([
-            gr.Tab("Upload Files", gr.Interface.Layout([
-                gr.File(label="Upload PDF or DOCX files", multiple=True),
-                gr.Button("Upload", onclick=upload_files),
-                gr.Textbox("Upload Status", default="No file uploaded yet", multiline=True)
-            ])),
-            gr.Tab("Query", gr.Interface.Layout([
-                gr.Textbox("Enter your query", label="Query Input"),
-                gr.Button("Search", onclick=process_and_query),
-                gr.Textbox("Query Response", default="No query processed yet", multiline=True)
-            ]))
-        ]),
-        outputs=gr.Textbox("Output", label="Output", default="Output will be shown here", multiline=True),
+        fn=None,
+        inputs=[
+            gr.Interface.Tab("Upload Files", upload_tab),
+            gr.Interface.Tab("Query", query_tab)
+        ],
+        outputs=gr.outputs.Textbox(label="Output", default="Output will be shown here", type="textbox"),
         live=True,
         capture_session=True
     ).launch()