EE21's picture
Update app.py
af90ec4
raw
history blame
7.19 kB
import streamlit as st
import re
from rouge import Rouge
from datasets import load_dataset
import PyPDF2
from extractive_summarization import summarize_with_textrank, summarize_with_lsa
from abstractive_summarization import summarize_with_bart_cnn, summarize_with_bart_ft, summarize_with_led, summarize_with_t5
from keyword_extraction import extract_keywords
from keyphrase_extraction import extract_sentences_with_obligations
#from blanc import BlancHelp
# Load in ToS
dataset = load_dataset("EE21/ToS-Summaries")
def extract_organization_name(text):
# A simple regex pattern to identify organization names. This pattern looks for capitalized words, possibly followed by "Inc.", "Ltd.", etc.
# This is a very basic pattern and might need to be adjusted based on the actual content of the documents.
pattern = r"\b[A-Z][a-zA-Z]*(?:\s[A-Z][a-zA-Z]*)*\s(?:Inc\.|Corporation|Corp\.|LLC|Ltd\.|Limited|Co\.|Company)?\b"
# Search for the pattern in the text
match = re.search(pattern, text)
if match:
return match.group()
else:
return "Unknown Organization"
# Extract titles or identifiers for the ToS
#tos_titles = [f"Document {i}" for i in range(len(dataset['train']))]
tos_titles = [extract_organization_name(doc['plan_text']) for doc in dataset['train']]
# Set page to wide mode
st.set_page_config(layout="wide")
# Function to handle file upload and return its content
def load_pdf(file):
pdf_reader = PyPDF2.PdfReader(file)
pdf_text = ""
for page_num in range(len(pdf_reader.pages)):
pdf_text += pdf_reader.pages[page_num].extract_text() or ""
return pdf_text
# Main app
def main():
st.title("Terms of Service Summarizer")
# Layout: 3 columns
col1, col2, col3 = st.columns([1, 3, 2], gap="large")
# Left column: Radio buttons for summarizer choice
with col1:
radio_options = ["Abstractive (T5)", "Abstractive (LED)", 'Abstractive (Fine-tuned BART)', "Abstractive (BART-large-CNN)", 'Extractive (TextRank)',
"Extractive (Latent Semantic Analysis)", 'Keyphrase Extraction (RAKE)', 'Keyword Extraction (RAKE)']
help_text = "Abstractive: Abstractive summarization generates a summary that may contain words not present in the original text. " \
"It uses a fine-tuned model on BART-large-CNN.<br>" \
"Extractive: Extractive summarization selects and extracts sentences or phrases directly from the original text to create a summary using the TextRank algorithm.<br>" \
"Keyword Extraction: Keyword extraction identifies and extracts important keywords or terms from the text using the Rake algorithm. " \
"These keywords can be used for various purposes such as content analysis and SEO.<br>" \
"Keyphrase Extraction: Keyphrase extraction is similar to keyword extraction but focuses on identifying multi-word phrases or expressions that are significant in the text using the Rake algorithm."
radio_selection = st.radio("Choose type of summarizer:", radio_options, help=help_text)
# Middle column: Text input and File uploader
with col2:
user_input = st.text_area("Enter your text here:")
uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
# Dropdown for selecting the document
tos_selection_index = st.selectbox("Select a Terms of Service Document", range(len(tos_titles)), format_func=lambda x: tos_titles[x])
if st.button("Summarize"):
if uploaded_file and user_input and tos_selection_index:
st.warning("Please provide either text input or a PDF file, not both.")
return
elif uploaded_file:
# Extract text from PDF
file_content = load_pdf(uploaded_file)
st.write("PDF uploaded successfully.")
elif user_input:
file_content = user_input
elif tos_selection_index is not None:
file_content = dataset['train'][tos_selection_index]['plain_text']
else:
st.warning("Please upload a PDF, enter some text, or select a document to summarize.")
return
# Perform extractive summarization
if radio_selection == "Extractive (TextRank)":
summary = summarize_with_textrank(file_content)
st.session_state.summary = summary
# Perform extractive summarization
if radio_selection == "Extractive (Latent Semantic Analysis)":
summary = summarize_with_lsa(file_content)
st.session_state.summary = summary
# Perform extractive summarization
if radio_selection == "Abstractive (Fine-tuned BART)":
summary = summarize_with_bart_ft(file_content)
st.session_state.summary = summary
# Perform extractive summarization
if radio_selection == "Abstractive (BART-large-CNN)":
summary = summarize_with_bart_cnn(file_content)
st.session_state.summary = summary
# Perform extractive summarization
if radio_selection == "Abstractive (T5)":
summary = summarize_with_t5(file_content)
st.session_state.summary = summary
# Perform extractive summarization
if radio_selection == "Abstractive (LED)":
summary = summarize_with_led(file_content)
st.session_state.summary = summary
# Perform Keyword Extraction
if radio_selection == "Keyword Extraction (RAKE)":
summary = extract_keywords(file_content)
st.session_state.summary = summary
# Perform Keyphrase Extraction
if radio_selection == "Keyphrase Extraction (RAKE)":
summary = extract_sentences_with_obligations(file_content)
st.session_state.summary = summary
# Right column: Displaying text after pressing 'Summarize'
with col3:
st.write("Summary:")
if 'summary' in st.session_state:
st.write(st.session_state.summary)
# Check if no PDF or text input is provided and a ToS document is selected
if not uploaded_file and not user_input and tos_selection_index is not None and 'summary' in dataset['train'][tos_selection_index]:
# Fetch the reference summary
reference_summary = dataset['train'][tos_selection_index]['summary']
# Calculate ROUGE scores
rouge = Rouge()
scores = rouge.get_scores(st.session_state.summary, reference_summary)
# Display ROUGE scores in a box
with st.container():
st.write("ROUGE Scores:")
st.write(f"ROUGE-1: {scores[0]['rouge-1']['f']:.4f}")
st.write(f"ROUGE-2: {scores[0]['rouge-2']['f']:.4f}")
st.write(f"ROUGE-L: {scores[0]['rouge-l']['f']:.4f}")
if __name__ == "__main__":
main()