Upload 3 files

app.py

@@ -2,40 +2,33 @@ import data_prep
 import model_predict
 import gradio as gr
 
-# Dictionary of model names and corresponding display names
 model_dict = {
-    "bert-base": "research-dump/bert-base-uncased_deletion_multiclass_complete_Final",
-    "bert-large": "research-dump/bert-large-uncased_deletion_multiclass_complete_final",
-    "roberta-base": "research-dump/roberta-base_deletion_multiclass_complete_final"
+    "BERT-Base": "research-dump/bert-base-uncased_deletion_multiclass_complete_Final",
+    "BERT-Large": "research-dump/bert-large-uncased_deletion_multiclass_complete_final",
+    "RoBERTa-Base": "research-dump/roberta-base_deletion_multiclass_complete_final",
+    "RoBERTa-Large": "research-dump/roberta-large_deletion_multiclass_complete_final"
 }
 
 def process_url(url, model_key):
-    # Get the actual model path from the model_dict
     model_name = model_dict[model_key]
-
-    # Process the text from the URL
     processed_text = data_prep.process_data(url)
-    
-    # Predict the labels and their probabilities
-    final_scores = model_predict.predict_text(processed_text, model_name)
-    
-    # Find the label with the highest probability
+    final_scores,highlighted_text = model_predict.predict_text(processed_text, model_name)
     highest_prob_label = max(final_scores, key=final_scores.get)
     highest_prob = final_scores[highest_prob_label]
-    
-    # Create progress bar style output for all labels
     progress_bars = {label: score for label, score in final_scores.items()}
     
-    return highest_prob_label, highest_prob, progress_bars
+    return processed_text, highest_prob_label, highest_prob, progress_bars #,highlighted_text
+
 
-# Define the interface for the Gradio app
 url_input = gr.Textbox(label="URL")
 model_name_input = gr.Dropdown(label="Model Name", choices=list(model_dict.keys()), value=list(model_dict.keys())[0])
 outputs = [
+    gr.Textbox(label="Processed Text"),
     gr.Textbox(label="Label with Highest Probability"),
     gr.Textbox(label="Probability"),
-    gr.JSON(label="All Labels and Probabilities")
+    gr.JSON(label="All Labels and Probabilities"),
+    #gr.HTML(label="Processed Text")
 ]
 
 demo = gr.Interface(fn=process_url, inputs=[url_input, model_name_input], outputs=outputs)
-demo.launch()
+demo.launch() #share=True)

data_prep.py

@@ -1,19 +1,18 @@
 import requests
 import pandas as pd
 from bs4 import BeautifulSoup
-import pandas as pd
-from datetime import datetime
-
 
 def extract_div_contents_from_url(url):
     response = requests.get(url)
     if response.status_code != 200:
+        print(f"Error: Received status code {response.status_code} for URL: {url}")
         return pd.DataFrame(columns=['title', 'text_url', 'deletion_discussion', 'label', 'confirmation', 'discussion', 'verdict'])
 
     soup = BeautifulSoup(response.content, 'html.parser')
-    divs = soup.find_all('div', class_='boilerplate afd vfd xfd-closed archived mw-archivedtalk')
-
-    # Extract the title fragment from the URL
+    div_classes = ['boilerplate afd vfd xfd-closed', 'boilerplate afd vfd xfd-closed archived mw-archivedtalk']
+    divs = []
+    for div_class in div_classes:
+        divs.extend(soup.find_all('div', class_=div_class))
     url_fragment = url.split('#')[-1].replace('_', ' ')
 
     data = []
@@ -43,9 +42,10 @@ def extract_div_contents_from_url(url):
                         title = title_tag.text
                         text_url = 'https://en.wikipedia.org' + title_tag['href']
 
-            if title != url_fragment:
-                continue  # Skip if the title does not match the URL fragment
-
+            if not title:
+                continue
+            if title.lower() != url_fragment.lower():
+                continue  
             deletion_discussion = div.prettify()
 
             # Extract label
@@ -58,11 +58,13 @@ def extract_div_contents_from_url(url):
 
             # Extract confirmation
             confirmation = ''
-            discussion_tag = div.find('dd').find('i')
+            discussion_tag = div.find('dd')
             if discussion_tag:
-                confirmation_b_tag = discussion_tag.find('b')
-                if confirmation_b_tag:
-                    confirmation = confirmation_b_tag.text.strip()
+                discussion_tag_i = discussion_tag.find('i')
+                if discussion_tag_i:
+                    confirmation_b_tag = discussion_tag_i.find('b')
+                    if confirmation_b_tag:
+                        confirmation = confirmation_b_tag.text.strip()
 
             # Split deletion_discussion into discussion and verdict
             parts = deletion_discussion.split('<div class="mw-heading mw-heading3">')
@@ -75,7 +77,8 @@ def extract_div_contents_from_url(url):
             continue
 
     df = pd.DataFrame(data, columns=['title', 'text_url', 'deletion_discussion', 'label', 'confirmation', 'verdict', 'discussion'])
-    df = df[['title','discussion','verdict','label']]
+    df = df[['title', 'discussion', 'verdict', 'label']]
+    print(f"DataFrame created with {len(df)} rows")
     return df
 
 
@@ -108,6 +111,7 @@ def process_html_to_plaintext(df):
     df = df[['title', 'discussion_cleaned', 'label']]
     return df
 
+
 import pysbd
 def split_text_into_sentences(text):
     seg = pysbd.Segmenter(language="en", clean=False)
@@ -117,8 +121,6 @@ def process_split_text_into_sentences(df):
     df['discussion_cleaned'] = df['discussion_cleaned'].apply(split_text_into_sentences)
     return df
 
-
-
 def process_data(url):
     df = extract_div_contents_from_url(url)
     df = process_discussion(df)
@@ -128,4 +130,5 @@ def process_data(url):
     if not df.empty:
         return df.at[0,'title']+ ' : '+df.at[0, 'discussion_cleaned']
     else:
-        return 'Empty DataFrame'
+        return 'Empty DataFrame'
+

model_predict.py

@@ -1,5 +1,39 @@
 #using pipeline to predict the input text
-from transformers import pipeline
+# from transformers import pipeline, AutoTokenizer
+# import torch
+
+# label_mapping = {
+#     'delete': [0, 'LABEL_0'],
+#     'keep': [1, 'LABEL_1'],
+#     'merge': [2, 'LABEL_2'],
+#     'no consensus': [3, 'LABEL_3'],
+#     'speedy keep': [4, 'LABEL_4'],
+#     'speedy delete': [5, 'LABEL_5'],
+#     'redirect': [6, 'LABEL_6'],
+#     'withdrawn': [7, 'LABEL_7']
+# }
+
+# def predict_text(text, model_name):
+#     tokenizer = AutoTokenizer.from_pretrained(model_name)
+#     model = pipeline("text-classification", model=model_name, return_all_scores=True)
+    
+#     # Tokenize and truncate the text
+#     tokens = tokenizer(text, truncation=True, max_length=512)
+#     truncated_text = tokenizer.decode(tokens['input_ids'], skip_special_tokens=True)
+    
+#     results = model(truncated_text)
+#     final_scores = {key: 0.0 for key in label_mapping}
+    
+#     for result in results[0]:
+#         for key, value in label_mapping.items():
+#             if result['label'] == value[1]:
+#                 final_scores[key] = result['score']
+#                 break
+    
+#     return final_scores
+
+
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
 
 label_mapping = {
@@ -14,15 +48,36 @@ label_mapping = {
 }
 
 def predict_text(text, model_name):
-    model = pipeline("text-classification", model=model_name, return_all_scores=True)
-    results = model(text)
-    final_scores = {key: 0.0 for key in label_mapping}
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, output_attentions=True)
+    
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    outputs = model(**inputs)
+    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
     
-    for result in results[0]:
+    final_scores = {key: 0.0 for key in label_mapping}
+    for i, score in enumerate(predictions[0]):
         for key, value in label_mapping.items():
-            if result['label'] == value[1]:
-                final_scores[key] = result['score']
+            if i == value[0]:
+                final_scores[key] = score.item()
                 break
     
-    return final_scores
-
+    # Calculate average attention
+    attentions = outputs.attentions
+    avg_attentions = torch.mean(torch.stack(attentions), dim=1)  # Average over all layers
+    avg_attentions = avg_attentions.mean(dim=1)[0]  # Average over heads
+    token_importance = avg_attentions.mean(dim=0)
+    
+    # Decode tokens and highlight important ones
+    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
+    highlighted_text = []
+    for token, importance in zip(tokens, token_importance):
+        if importance > token_importance.mean():
+            highlighted_text.append(f"<b>{token}</b>") #
+        else:
+            highlighted_text.append(token)
+    
+    highlighted_text = " ".join(highlighted_text)
+    highlighted_text = highlighted_text.replace("##", "") 
+    
+    return final_scores, highlighted_text