shaima21's picture
Create app.py
bcbe6f2 verified
import os
import streamlit as st
from transformers import pipeline
import torch
from PyPDF2 import PdfReader
# Disable tokenizers parallelism
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Setup for the model
device = 0 if torch.cuda.is_available() else -1
summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", device=device)
def split_text(text, max_chunk_size=512):
words = text.split()
for i in range(0, len(words), max_chunk_size):
yield " ".join(words[i:i + max_chunk_size])
def extract_text_from_pdf(pdf_file):
reader = PdfReader(pdf_file)
text = ""
for page_num in range(len(reader.pages)):
page = reader.pages[page_num]
text += page.extract_text()
return text
def summarize_text(text, summarizer):
chunks = list(split_text(text))
summaries = []
for chunk in chunks:
input_length = len(chunk.split())
max_summary_length = max(10, int(input_length * 0.6))
min_summary_length = max(5, int(input_length * 0.2))
result = summarizer(chunk, max_length=max_summary_length, min_length=min_summary_length, do_sample=False)
summaries.append(result[0]['summary_text'])
return " ".join(summaries)
def extract_and_summarize_page_by_page(pdf_file, summarizer):
reader = PdfReader(pdf_file)
summaries = []
for page_num in range(len(reader.pages)):
page = reader.pages[page_num]
text = page.extract_text()
if text:
page_summary = summarize_text(text, summarizer)
summaries.append(page_summary)
else:
summaries.append(f"Page {page_num + 1}: No extractable text found.")
return summaries
# Streamlit interface
st.subheader("Generate PDF Summary")
pdf_file = st.file_uploader("Upload a PDF", type=["pdf"])
if pdf_file:
text = extract_text_from_pdf(pdf_file)
if len(text) > 0:
summaries = extract_and_summarize_page_by_page(pdf_file, summarizer)
st.subheader("Summary")
for i, summary in enumerate(summaries, 1):
st.write(f"### Page {i}\n{summary}\n")
else:
st.warning("No extractable text found in the PDF.")