File size: 3,289 Bytes
a742457 18085b7 9d6cd25 e48c184 9d6cd25 18085b7 a742457 18085b7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
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)}")
|