Update app.py

app.py

@@ -2,11 +2,15 @@ import streamlit as st
 from transformers import pipeline
 import PyPDF2
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from sentence_transformers import SentenceTransformer, util
 
-# Load Hugging Face Question Answering model
-qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
+# Load the Question Answering Model
+qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
 
-# Function to extract text from PDF
+# Load Embeddings Model for Better Context Matching
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+# Function to Extract Text from PDF
 def extract_text_from_pdf(pdf_path):
     with open(pdf_path, "rb") as f:
         pdf_reader = PyPDF2.PdfReader(f)
@@ -15,31 +19,31 @@ def extract_text_from_pdf(pdf_path):
             text += page.extract_text() + "\n"
     return text
 
-# Function to split text into smaller chunks
+# Function to Split Text into Chunks
 def split_text(text):
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
     chunks = text_splitter.split_text(text)
     return chunks
 
-# Function to find the most relevant chunk for a question
-def find_relevant_chunk(question, chunks):
-    best_chunk = ""
-    best_score = 0
-    for chunk in chunks:
-        response = qa_pipeline(question=question, context=chunk)
-        score = response['score']
-        if score > best_score:
-            best_score = score
-            best_chunk = chunk
-    return best_chunk
+# Function to Find the Most Relevant Chunk Using Embeddings
+def find_best_chunk(question, chunks):
+    question_embedding = embedding_model.encode(question, convert_to_tensor=True)
+    chunk_embeddings = [embedding_model.encode(chunk, convert_to_tensor=True) for chunk in chunks]
+
+    # Compute similarity between question and each chunk
+    similarities = [util.pytorch_cos_sim(question_embedding, chunk_emb).item() for chunk_emb in chunk_embeddings]
+
+    # Find the most relevant chunk
+    best_chunk_index = similarities.index(max(similarities))
+    return chunks[best_chunk_index]
 
 # Streamlit UI
 st.title("Chat with AWS Restart PDF")
 
-# Use the uploaded PDF file
-pdf_path = "AWS restart program information.docx.pdf"  # Update with your file name
+# Load and Process PDF
+pdf_path = "AWS restart program information.docx.pdf"
 pdf_text = extract_text_from_pdf(pdf_path)
-chunks = split_text(pdf_text)  # Split the text into chunks
+chunks = split_text(pdf_text)
 
 st.write("✅ PDF Loaded Successfully!")
 
@@ -47,6 +51,6 @@ st.write("✅ PDF Loaded Successfully!")
 question = st.text_input("Ask a question about AWS Restart program:")
 
 if st.button("Get Answer") and question:
-    relevant_chunk = find_relevant_chunk(question, chunks)  # Get the best chunk
-    response = qa_pipeline(question=question, context=relevant_chunk)  # Ask model on best chunk
+    relevant_chunk = find_best_chunk(question, chunks)  # Retrieve the best chunk
+    response = qa_pipeline(question=question, context=relevant_chunk)  # Ask the model
     st.write("Answer:", response['answer'])