Update app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+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
@@ -6,6 +7,13 @@ 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")
+
+# Extract titles or identifiers for the ToS
+tos_titles = [f"Document {i}" for i in range(len(dataset['train']))]
+
+
 # Set page to wide mode
 st.set_page_config(layout="wide")
 
@@ -42,6 +50,10 @@ def main():
     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 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:
                 st.warning("Please provide either text input or a PDF file, not both.")
@@ -52,8 +64,10 @@ def main():
                 st.write("PDF uploaded successfully.")
             elif user_input:
                 file_content = user_input
+            elif tos_selection_index is not None:  # Assuming tos_selection_index won't be None
+                file_content = dataset['train'][tos_selection_index]['plain_text']
             else:
-                st.warning("Please upload a PDF or enter some text to summarize.")
+                st.warning("Please upload a PDF, enter some text, or select a document to summarize.")
                 return
 
             # Perform extractive summarization