Update app.py

app.py

@@ -1,6 +1,5 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
-from sentence_transformers import SentenceTransformer
+from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
 from pydantic import BaseModel, Field
 from typing import List
 from datetime import datetime
@@ -8,20 +7,24 @@ 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
 import time
 
-# === Load summarization model ===
+# === Summarization model ===
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
-# === Load QA pipeline ===
-qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
+# === QA model ===
+qa_tokenizer = AutoTokenizer.from_pretrained("deepset/roberta-base-squad2")
+qa_model = AutoModelForQuestionAnswering.from_pretrained("deepset/roberta-base-squad2")
+qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)
 
-# === Load BGE-small embedding model ===
+# === Embedding model ===
+from sentence_transformers import SentenceTransformer
+from langchain.embeddings import HuggingFaceEmbeddings
 embedding_model = SentenceTransformer("BAAI/bge-small-en-v1.5")
+embedding_function = HuggingFaceEmbeddings(model=embedding_model)
 
 # === Data models ===
 class KeyPoint(BaseModel):
@@ -39,9 +42,20 @@ def extract_text_from_pdf(pdf_file):
     return "".join(page.extract_text() for page in pdf_reader.pages)
 
 def analyze_text_structured(text):
-    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)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.split_text(text)
+
+    summaries = []
+    for chunk in chunks:
+        try:
+            result = summarizer(chunk, max_length=200, min_length=50, do_sample=False)
+            summaries.append(result[0]["summary_text"])
+        except Exception:
+            summaries.append("")
+
+    full_summary = " ".join(summaries)
+    key_points = [KeyPoint(point=line.strip()) for line in full_summary.split(". ") if line.strip()]
+    return DocumentAnalysis(summary=Summary(summary=full_summary), key_points=key_points)
 
 def json_to_text(analysis):
     text_output = "=== Summary ===\n" + f"{analysis.summary.summary}\n\n"
@@ -74,9 +88,9 @@ def create_word_report(analysis):
     return docx_bytes.getvalue()
 
 # === Streamlit UI ===
-st.set_page_config(page_title="Chat With PDF (BART + BGE + RoBERTa)", page_icon="📄")
+st.set_page_config(page_title="Chat With PDF (BART + BGE)", page_icon="📄")
 st.title("📄 Chat With PDF")
-st.caption("Summarize and Chat with Documents using facebook/bart-large-cnn + BGE-small + RoBERTa QA")
+st.caption("Summarize and Chat with Documents using facebook/bart-large-cnn + BGE-small Embeddings + RoBERTa QA")
 
 for key in ["current_file", "pdf_summary", "analysis_time", "pdf_report", "word_report", "vectorstore", "messages"]:
     if key not in st.session_state:
@@ -98,12 +112,10 @@ if uploaded_file is not None:
             analysis = analyze_text_structured(text)
             st.session_state.pdf_summary = analysis
 
-            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-            chunks = text_splitter.split_text(text)
+            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+            chunks = splitter.split_text(text)
             docs = [LCDocument(page_content=chunk) for chunk in chunks]
-
-            vectors = embedding_model.encode([doc.page_content for doc in docs])
-            st.session_state.vectorstore = FAISS.from_embeddings(docs, vectors)
+            st.session_state.vectorstore = FAISS.from_documents(docs, embedding_function)
 
             st.session_state.pdf_report = create_pdf_report(analysis)
             st.session_state.word_report = create_word_report(analysis)
@@ -143,12 +155,12 @@ 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})
 
 if st.session_state.analysis_time is not None:
     st.markdown(
-        f'<div style="text-align:center; margin-top:2rem; color:gray;">Analysis Time: {st.session_state.analysis_time:.1f}s | Embedding: BGE Small v1.5</div>',
+        f'<div style="text-align:center; margin-top:2rem; color:gray;">Analysis Time: {st.session_state.analysis_time:.1f}s | Embedding: BGE-small v1.5 | QA: RoBERTa-SQuAD2</div>',
         unsafe_allow_html=True
     )