Update app.py

app.py

@@ -22,15 +22,20 @@ id2label = {
 
 # Color map for entities
 label_colors = {
-    "B-STEREO": "#FFCDD2",
-    "I-STEREO": "#E57373",
-    "B-GEN": "#C8E6C9",
-    "I-GEN": "#81C784",
-    "B-UNFAIR": "#BBDEFB",
-    "I-UNFAIR": "#64B5F6",
-    "O": "#FFFFFF"  # Default for no label
+    "STEREO": "rgba(255, 0, 0, 0.3)",  # Red
+    "GEN": "rgba(0, 0, 255, 0.3)",     # Blue
+    "UNFAIR": "rgba(0, 255, 0, 0.3)"   # Green
 }
 
+# Helper to wrap a token in a span with color
+def wrap_token_with_color(token, labels):
+    # Build nested highlights
+    style = "position: relative;"
+    for label in labels:
+        if label != "O":
+            style += f"background: {label_colors[label]};"
+    return f"<span style='{style}'>{token}</span>"
+
 # Predict function
 def predict_ner_tags(sentence):
     inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=128)
@@ -43,19 +48,34 @@ def predict_ner_tags(sentence):
         probabilities = torch.sigmoid(logits)
         predicted_labels = (probabilities > 0.5).int()  # Threshold
 
-    result = []
     tokens = tokenizer.convert_ids_to_tokens(input_ids[0])
     highlighted_sentence = ""
+    prev_labels = []
+
     for i, token in enumerate(tokens):
         if token not in tokenizer.all_special_tokens:
+            # Extract the labels for this token
             label_indices = (predicted_labels[0][i] == 1).nonzero(as_tuple=False).squeeze(-1)
-            labels = [id2label[idx.item()] for idx in label_indices] if label_indices.numel() > 0 else ['O']
-            # Get the most prominent label for coloring (arbitrary choice for multiple labels)
-            primary_label = labels[0] if labels else "O"
-            color = label_colors.get(primary_label, "#FFFFFF")
-            highlighted_sentence += f"<span style='background-color:{color}'>{token}</span> "
+            labels = [id2label[idx.item()][2:] for idx in label_indices] if label_indices.numel() > 0 else ['O']
+            
+            # Check if labels are the same as the previous token (for seamless highlighting)
+            if labels != prev_labels:
+                if prev_labels:  # Close the previous span if needed
+                    highlighted_sentence += "</span>"
+
+                # Start a new span
+                if labels != ["O"]:
+                    highlighted_sentence += f"<span style='background: linear-gradient({', '.join([label_colors[label] for label in labels])})'>"
+            
+            # Add the token to the span
+            highlighted_sentence += token.replace("##", "")
+            prev_labels = labels
+
+    # Close any open spans
+    if prev_labels and prev_labels != ["O"]:
+        highlighted_sentence += "</span>"
 
-    return highlighted_sentence.strip()
+    return highlighted_sentence
 
 # Gradio Interface
 iface = gr.Interface(