Update app.py

app.py

@@ -2,11 +2,13 @@ import torch
 from transformers import BertTokenizerFast, BertForTokenClassification
 import gradio as gr
 
+# Load tokenizer and model
 tokenizer = BertTokenizerFast.from_pretrained('bert-base-uncased')
 model = BertForTokenClassification.from_pretrained('maximuspowers/bias-detection-ner')
-model.eval()
-model.to('cuda' if torch.cuda.is_available() else 'cpu')
+model.eval()  # Set the model to evaluation mode
+model.to('cuda' if torch.cuda.is_available() else 'cpu')  # Move model to appropriate device
 
+# Define label mappings
 id2label = {
     0: 'O',
     1: 'B-STEREO',
@@ -18,10 +20,12 @@ id2label = {
 }
 
 def predict_ner_tags(sentence):
+    # Tokenize the sentence
     inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=128)
     input_ids = inputs['input_ids'].to(model.device)
     attention_mask = inputs['attention_mask'].to(model.device)
 
+    # Predict using the model
     with torch.no_grad():
         outputs = model(input_ids=input_ids, attention_mask=attention_mask)
         logits = outputs.logits
@@ -39,19 +43,31 @@ def predict_ner_tags(sentence):
     return result
 
 def format_output(result):
-    formatted_output = ""
+    # Create HTML content with formatted output
+    formatted_output = "<div style='font-family: Arial;'>"
     for token, labels in result:
-        formatted_output += f"{token}: {', '.join(labels)}\n"
+        styles = []
+        if "B-STEREO" in labels or "I-STEREO" in labels:
+            styles.append("border-bottom: 2px solid blue;")
+        if "B-GEN" in labels or "I-GEN" in labels:
+            styles.append("background-color: green; color: white;")
+        if "B-UNFAIR" in labels or "I-UNFAIR" in labels:
+            styles.append("border: 2px dashed red;")
+        
+        style_string = " ".join(styles) if styles else ""
+        formatted_output += f"<span style='{style_string} padding: 3px; margin: 2px;'>{token}</span> "
+    formatted_output += "</div>"
     return formatted_output
 
 iface = gr.Interface(
     fn=predict_ner_tags,
     inputs="text",
-    outputs="text",
+    outputs=gr.outputs.HTML(label="Output"),
     title="Named Entity Recognition with BERT",
-    description="Enter a sentence to predict NER tags using BERT model trained for multi-label classification.",
+    description="Enter a sentence to predict NER tags using a BERT model trained for multi-label classification. Different styles represent different entity types.",
     examples=["Tall men are so clumsy."],
-    allow_flagging="never"
+    allow_flagging="never",
+    theme="default"
 )
 
 if __name__ == "__main__":