Update app.py

app.py

@@ -7,28 +7,35 @@ 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 transformers import AutoModelForTokenClassification, AutoTokenizer
+import torch
 #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"
+model_name = "dbmdz/bert-large-cased-finetuned-conll03-english"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForTokenClassification.from_pretrained(model_name)
+
+def extract_organization_names(text):
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(**inputs)
+
+    predictions = torch.argmax(outputs.logits, dim=2)
+    entities = [tokenizer.convert_ids_to_tokens(inputs.input_ids[0][idx]) for idx, pred in enumerate(predictions[0]) if model.config.id2label[pred.item()] == 'B-ORG']
+    return " ".join(entities)
+
+# Apply this function to your dataset
+tos_titles = [extract_organization_names(doc['plain_text']) for doc in dataset['train']]
+
 
 # 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['plain_text']) for doc in dataset['train']]
+
 
         
 # Set page to wide mode