|
|
import streamlit as st |
|
|
import os |
|
|
import time |
|
|
from datetime import datetime |
|
|
from datasets import load_dataset |
|
|
from PyPDF2 import PdfReader |
|
|
from transformers import pipeline |
|
|
from sentence_transformers import SentenceTransformer |
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter |
|
|
import faiss |
|
|
import numpy as np |
|
|
import os |
|
|
os.makedirs("/app/cache", exist_ok=True) |
|
|
os.makedirs("/.streamlit", exist_ok=True) |
|
|
|
|
|
|
|
|
|
|
|
st.title("Document Summarization") |
|
|
|
|
|
@st.cache_resource |
|
|
def load_models(): |
|
|
embedder = SentenceTransformer("all-MiniLM-L6-v2") |
|
|
summarizer = pipeline("summarization", model="facebook/bart-large-cnn") |
|
|
return embedder, summarizer |
|
|
|
|
|
embedder, summarizer = load_models() |
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50, length_function=len) |
|
|
|
|
|
def process_file(uploaded_file): |
|
|
if uploaded_file.name.endswith('.pdf'): |
|
|
text = "" |
|
|
reader = PdfReader(uploaded_file) |
|
|
for page in reader.pages: |
|
|
text += page.extract_text() or "" |
|
|
elif uploaded_file.name.endswith(('.txt', '.md')): |
|
|
text = uploaded_file.read().decode("utf-8") |
|
|
else: |
|
|
raise ValueError("Unsupported file format. Use PDF, TXT, or Markdown.") |
|
|
return text |
|
|
|
|
|
st.header("View Summaries of 3 Dataset Documents") |
|
|
dataset = load_dataset("cnn_dailymail", "3.0.0") |
|
|
train_samples = dataset["train"].select(range(3)) |
|
|
|
|
|
for i, sample in enumerate(train_samples, 1): |
|
|
context = sample["article"] |
|
|
reference = sample["highlights"] |
|
|
with st.expander(f"Document {i}"): |
|
|
start_time = time.time() |
|
|
summary = summarizer(context, max_length=150)[0]['summary_text'] |
|
|
latency = time.time() - start_time |
|
|
|
|
|
st.subheader("Generated Summary") |
|
|
st.write(summary) |
|
|
st.subheader("Reference Summary") |
|
|
st.write(reference) |
|
|
st.text(f"Latency: {latency:.2f}s | Input tokens: {len(context.split())} | Output tokens: {len(summary.split())}") |
|
|
|
|
|
st.header("Upload Your Own Document") |
|
|
input_text = st.text_area("Or paste your text here") |
|
|
uploaded_file = st.file_uploader("Upload a PDF, TXT, or Markdown file") |
|
|
|
|
|
if st.button("Summarize"): |
|
|
try: |
|
|
if uploaded_file: |
|
|
text = process_file(uploaded_file) |
|
|
elif input_text.strip(): |
|
|
text = input_text.strip() |
|
|
else: |
|
|
st.warning("Please provide text input or upload a file.") |
|
|
st.stop() |
|
|
|
|
|
chunks = text_splitter.split_text(text) |
|
|
chunk_embeddings = embedder.encode(chunks) |
|
|
index = faiss.IndexFlatL2(chunk_embeddings.shape[1]) |
|
|
index.add(chunk_embeddings) |
|
|
|
|
|
query_embedding = embedder.encode(["Summarize this document"]) |
|
|
_, indices = index.search(query_embedding, 3) |
|
|
retrieved_chunks = [chunks[i] for i in indices[0]] |
|
|
context = " ".join(retrieved_chunks) |
|
|
|
|
|
start_time = time.time() |
|
|
summary = summarizer(context, max_length=150)[0]['summary_text'] |
|
|
latency = time.time() - start_time |
|
|
|
|
|
st.subheader("Generated Summary") |
|
|
st.write(summary) |
|
|
st.subheader("Context") |
|
|
st.write(context) |
|
|
st.text(f"Latency: {latency:.2f}s | Input tokens: {len(context.split())} | Output tokens: {len(summary.split())}") |
|
|
|
|
|
except Exception as e: |
|
|
st.error(f"Error: {str(e)}") |
|
|
|
