Update app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+import re
 from rouge import Rouge
 from datasets import load_dataset
 import PyPDF2
@@ -12,9 +13,24 @@ from keyphrase_extraction import extract_sentences_with_obligations
 # 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 = [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")