Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
+from sentence_transformers import SentenceTransformer
 from pydantic import BaseModel, Field
 from typing import List
 from datetime import datetime
@@ -7,35 +8,21 @@ import PyPDF2
 from fpdf import FPDF
 from docx import Document
 import io
+import numpy as np
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
 from langchain_core.documents import Document as LCDocument
-from langchain_core.embeddings import Embeddings
-from sentence_transformers import SentenceTransformer
 import time
 
 # === Load summarization model ===
-tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
-model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
 # === Load QA pipeline ===
-qa_pipeline = pipeline("question-answering", model="facebook/bart-large-cnn", tokenizer=tokenizer)
+qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
 
-# === Load SentenceTransformer embedding model ===
+# === Load BGE-small embedding model ===
 embedding_model = SentenceTransformer("BAAI/bge-small-en-v1.5")
 
-class CustomSentenceTransformer(Embeddings):
-    def __init__(self, model):
-        self.model = model
-
-    def embed_documents(self, texts):
-        return self.model.encode(texts, show_progress_bar=False).tolist()
-
-    def embed_query(self, text):
-        return self.model.encode(text, show_progress_bar=False).tolist()
-
-embedding_function = CustomSentenceTransformer(embedding_model)
-
 # === Data models ===
 class KeyPoint(BaseModel):
     point: str = Field(description="A key point extracted from the document.")
@@ -52,14 +39,9 @@ def extract_text_from_pdf(pdf_file):
     return "".join(page.extract_text() for page in pdf_reader.pages)
 
 def analyze_text_structured(text):
-    inputs = tokenizer([text], max_length=1024, truncation=True, return_tensors="pt")
-    summary_ids = model.generate(
-        inputs["input_ids"], num_beams=4, length_penalty=2.0,
-        max_length=200, min_length=50, early_stopping=True
-    )
-    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
-    key_points = [KeyPoint(point=line.strip()) for line in summary.split(". ") if line.strip()]
-    return DocumentAnalysis(summary=Summary(summary=summary), key_points=key_points)
+    result = summarizer(text, max_length=200, min_length=50, do_sample=False)[0]["summary_text"]
+    key_points = [KeyPoint(point=line.strip()) for line in result.split(". ") if line.strip()]
+    return DocumentAnalysis(summary=Summary(summary=result), key_points=key_points)
 
 def json_to_text(analysis):
     text_output = "=== Summary ===\n" + f"{analysis.summary.summary}\n\n"
@@ -92,9 +74,9 @@ def create_word_report(analysis):
     return docx_bytes.getvalue()
 
 # === Streamlit UI ===
-st.set_page_config(page_title="Chat With PDF (BART + BGE)", page_icon="📄")
+st.set_page_config(page_title="Chat With PDF (BART + BGE + RoBERTa)", page_icon="📄")
 st.title("📄 Chat With PDF")
-st.caption("Summarize and Chat with Documents using facebook/bart-large-cnn + BGE Small Embeddings")
+st.caption("Summarize and Chat with Documents using facebook/bart-large-cnn + BGE-small + RoBERTa QA")
 
 for key in ["current_file", "pdf_summary", "analysis_time", "pdf_report", "word_report", "vectorstore", "messages"]:
     if key not in st.session_state:
@@ -120,7 +102,8 @@ if uploaded_file is not None:
             chunks = text_splitter.split_text(text)
             docs = [LCDocument(page_content=chunk) for chunk in chunks]
 
-            st.session_state.vectorstore = FAISS.from_documents(docs, embedding_function)
+            vectors = embedding_model.encode([doc.page_content for doc in docs])
+            st.session_state.vectorstore = FAISS.from_embeddings(docs, vectors)
 
             st.session_state.pdf_report = create_pdf_report(analysis)
             st.session_state.word_report = create_word_report(analysis)
@@ -160,7 +143,7 @@ if st.session_state.vectorstore is not None:
             with st.spinner("Searching..."):
                 docs = st.session_state.vectorstore.similarity_search(prompt, k=3)
                 context = "\n".join([doc.page_content for doc in docs])
-                answer = qa_pipeline({"question": prompt, "context": context})["answer"]
+                answer = qa_pipeline(question=prompt, context=context)["answer"]
                 st.markdown(answer)
         st.session_state.messages.append({"role": "assistant", "content": answer})