Update app.py

app.py

@@ -1,101 +1,159 @@
-# app.py
-
-from fastapi import FastAPI, UploadFile, File
-from pydantic import BaseModel
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from pydantic import BaseModel, Field
 from typing import List
-import fitz  # PyMuPDF
-from transformers import pipeline
-from sentence_transformers import SentenceTransformer
-from langchain.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from datetime import datetime
+import PyPDF2
+from fpdf import FPDF
+from docx import Document
+import io
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain.schema import Document
-from langchain.chains.question_answering import load_qa_chain
-from langchain.llms import HuggingFacePipeline
-from langchain_core.documents import Document as LangchainDocument
-
-# --- Init FastAPI ---
-app = FastAPI()
+from langchain_community.vectorstores import FAISS
+from langchain_core.documents import Document as LCDocument
+from langchain.embeddings import HuggingFaceEmbeddings
+import time
 
-# --- Summarizer ---
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
+# === Load summarization model ===
+tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
+model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
 
-# --- Question Answering ---
-qa_pipe = pipeline("question-answering", model="deepset/roberta-base-squad2")
+# === Load QA pipeline ===
+qa_pipeline = pipeline("question-answering", model="facebook/bart-large-cnn", tokenizer=tokenizer)
 
-# --- Embedding model ---
-embedding_model = HuggingFaceBgeEmbeddings(model_name="BAAI/bge-small-en-v1.5")
+# === Setup BGE Embedding model ===
+embedding_model_name = "BAAI/bge-large-en-v1.5"
+embedding_function = HuggingFaceEmbeddings(model_name=embedding_model_name)
 
-# --- Text Splitter ---
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
+# === Data models ===
+class KeyPoint(BaseModel):
+    point: str = Field(description="A key point extracted from the document.")
 
-# --- Pydantic schemas ---
 class Summary(BaseModel):
-    summary: str
-
-class KeyPoint(BaseModel):
-    point: str
+    summary: str = Field(description="A brief summary of the document content.")
 
 class DocumentAnalysis(BaseModel):
-    summary: Summary
     key_points: List[KeyPoint]
+    summary: Summary
 
-class QARequest(BaseModel):
-    question: str
-    context: str
-
-class QAResponse(BaseModel):
-    answer: str
-
-# --- PDF Text Extractor ---
-def extract_text_from_pdf(pdf_file: UploadFile) -> str:
-    text = ""
-    with fitz.open(stream=pdf_file.file.read(), filetype="pdf") as doc:
-        for page in doc:
-            text += page.get_text()
-    return text
-
-# --- Analyze Text (summarization) ---
-def analyze_text_structured(text: str) -> DocumentAnalysis:
-    chunks = text_splitter.split_text(text)
-    summaries = []
-    for chunk in chunks:
-        result = summarizer(chunk, max_length=200, min_length=50, do_sample=False)
-        if result:
-            summaries.append(result[0]["summary_text"])
-
-    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)
-
-# --- Question Answering ---
-def answer_question(question: str, context: str) -> str:
-    result = qa_pipe(question=question, context=context)
-    return result["answer"]
-
-# --- PDF Upload + Analysis Route ---
-@app.post("/analyze-pdf", response_model=DocumentAnalysis)
-async def analyze_pdf(file: UploadFile = File(...)):
-    text = extract_text_from_pdf(file)
-    analysis = analyze_text_structured(text)
-    return analysis
-
-# --- Question Answering Route ---
-@app.post("/qa", response_model=QAResponse)
-async def ask_question(qa_request: QARequest):
-    answer = answer_question(qa_request.question, qa_request.context)
-    return QAResponse(answer=answer)
-
-# --- Embedding Search (FAISS) Demo ---
-@app.post("/search-chunks")
-async def search_chunks(file: UploadFile = File(...), query: str = ""):
-    text = extract_text_from_pdf(file)
-    chunks = text_splitter.split_text(text)
-    documents = [LangchainDocument(page_content=chunk) for chunk in chunks]
-
-    # Create FAISS vector store
-    db = FAISS.from_documents(documents, embedding_model)
-
-    # Similarity search
-    results = db.similarity_search(query, k=3)
-    return {"results": [doc.page_content for doc in results]}
+def extract_text_from_pdf(pdf_file):
+    pdf_reader = PyPDF2.PdfReader(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)
+
+def json_to_text(analysis):
+    text_output = "=== Summary ===\n" + f"{analysis.summary.summary}\n\n"
+    text_output += "=== Key Points ===\n"
+    for i, key_point in enumerate(analysis.key_points, start=1):
+        text_output += f"{i}. {key_point.point}\n"
+    return text_output
+
+def create_pdf_report(analysis):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font('Helvetica', '', 12)
+    pdf.cell(200, 10, txt="PDF Analysis Report", ln=True, align='C')
+    pdf.cell(200, 10, txt=f"Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=True, align='C')
+    pdf.multi_cell(0, 10, txt=json_to_text(analysis))
+    pdf_bytes = io.BytesIO()
+    pdf.output(pdf_bytes, dest='S')
+    pdf_bytes.seek(0)
+    return pdf_bytes.getvalue()
+
+def create_word_report(analysis):
+    doc = Document()
+    doc.add_heading('PDF Analysis Report', 0)
+    doc.add_paragraph(f'Generated on: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}')
+    doc.add_heading('Analysis', level=1)
+    doc.add_paragraph(json_to_text(analysis))
+    docx_bytes = io.BytesIO()
+    doc.save(docx_bytes)
+    docx_bytes.seek(0)
+    return docx_bytes.getvalue()
+
+# === Streamlit UI ===
+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 Embeddings")
+
+for key in ["current_file", "pdf_summary", "analysis_time", "pdf_report", "word_report", "vectorstore", "messages"]:
+    if key not in st.session_state:
+        st.session_state[key] = None if key != "messages" else []
+
+uploaded_file = st.file_uploader("Upload a PDF file", type="pdf")
+
+if uploaded_file is not None:
+    if st.session_state.current_file != uploaded_file.name:
+        st.session_state.current_file = uploaded_file.name
+        for key in ["pdf_summary", "pdf_report", "word_report", "vectorstore", "messages"]:
+            st.session_state[key] = None if key != "messages" else []
+
+    text = extract_text_from_pdf(uploaded_file)
+
+    if st.button("Analyze Text"):
+        start_time = time.time()
+        with st.spinner("Analyzing..."):
+            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)
+            docs = [LCDocument(page_content=chunk) for chunk in chunks]
+
+            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)
+
+        st.session_state.analysis_time = time.time() - start_time
+        st.subheader("Analysis Results")
+        st.text(json_to_text(analysis))
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.download_button(
+                label="Download PDF Report",
+                data=st.session_state.pdf_report,
+                file_name="analysis_report.pdf",
+                mime="application/pdf"
+            )
+        with col2:
+            st.download_button(
+                label="Download Word Report",
+                data=st.session_state.word_report,
+                file_name="analysis_report.docx",
+                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+            )
+
+if st.session_state.vectorstore is not None:
+    st.subheader("Chat with the Document")
+
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    if prompt := st.chat_input("Ask a question about the document"):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        with st.chat_message("assistant"):
+            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"]
+                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 Large v1.5</div>',
+        unsafe_allow_html=True
+    )