Update app.py

app.py

@@ -32,7 +32,7 @@ def wrap_token_with_color(token, labels):
     # Build nested highlights
     style = "position: relative;"
     for label in labels:
-        if label != "O":
+        if label != "O" and label in label_colors:
             style += f"background: {label_colors[label]};"
     return f"<span style='{style}'>{token}</span>"
 
@@ -56,7 +56,10 @@ def predict_ner_tags(sentence):
         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()][2:] for idx in label_indices] if label_indices.numel() > 0 else ['O']
+            labels = [id2label[idx.item()][2:] for idx in label_indices if idx.item() in id2label]  # Safe lookup
+            
+            if not labels:  # Handle empty labels gracefully
+                labels = ["O"]
             
             # Check if labels are the same as the previous token (for seamless highlighting)
             if labels != prev_labels:
@@ -65,7 +68,9 @@ def predict_ner_tags(sentence):
 
                 # Start a new span
                 if labels != ["O"]:
-                    highlighted_sentence += f"<span style='background: linear-gradient({', '.join([label_colors[label] for label in labels])})'>"
+                    highlight_colors = [label_colors[label] for label in labels if label in label_colors]
+                    if highlight_colors:  # Only create gradient if valid colors exist
+                        highlighted_sentence += f"<span style='background: linear-gradient({', '.join(highlight_colors)});'>"
             
             # Add the token to the span
             highlighted_sentence += token.replace("##", "")
@@ -90,4 +95,4 @@ iface = gr.Interface(
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch(share=True)