Update app.py

app.py

@@ -2,16 +2,13 @@ import streamlit as st
 import pdfplumber
 import faiss
 import numpy as np
-import torch
 from sentence_transformers import SentenceTransformer
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import pipeline
 
-# Load Flan-T5 Model for Detailed Answers
-model_name = "google/flan-t5-small"  # Smaller version that loads faster
-
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+# Load Extractive QA Model (Like ChatPDF)
+model_name = "deepset/roberta-base-squad2"
+qa_pipeline = pipeline("question-answering", model=model_name)
 
 # Load Sentence Embeddings Model
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
@@ -48,16 +45,10 @@ def find_best_chunk(question, index, chunks, embeddings):
     _, closest_idx = index.search(np.array(question_embedding), 1)
     return chunks[closest_idx[0][0]]
 
-# Function to Generate a Long, Detailed Answer
+# Function to Extract the Best Answer
 def get_answer(question, context):
-    input_text = f"Question: {question}\nContext: {context}\nAnswer:"
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
-    
-    # Generate response
-    with torch.no_grad():
-        output = model.generate(**inputs, max_length=300, temperature=0.7)
-    
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    response = qa_pipeline(question=question, context=context)
+    return response['answer']  # Returns extracted answer (ChatPDF-like behavior)
 
 # Streamlit UI
 st.title("Chat with AWS Restart PDF (Like ChatPDF)")