Update app.py

app.py

@@ -86,32 +86,15 @@ label_column_name = "ner_tags"
 
 # In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the
 # unique labels.
-def get_label_list(labels):
-    unique_labels = set()
-    for label in labels:
-        unique_labels = unique_labels | set(label)
-    label_list = list(unique_labels)
-    label_list.sort()
-    return label_list
-
-if isinstance(features[label_column_name].feature, ClassLabel):
-    label_list = features[label_column_name].feature.names
-    # No need to convert the labels since they are already ints.
-    id2label = {k: v for k,v in enumerate(label_list)}
-    label2id = {v: k for k,v in enumerate(label_list)}
-else:
-    label_list = get_label_list(dataset["test"][label_column_name])
-    id2label = {k: v for k,v in enumerate(label_list)}
-    label2id = {v: k for k,v in enumerate(label_list)}
-num_labels = len(label_list)
-
-def get_label_list(labels):
-    unique_labels = set()
-    for label in labels:
-        unique_labels = unique_labels | set(label)
-    label_list = list(unique_labels)
-    label_list.sort()
-    return label_list
+labels = dataset.features["ner_tags"].feature.names
+id2label = {v: k for v, k in enumerate(labels)}
+
+label2color = {
+    "question": "blue",
+    "answer": "green",
+    "header": "orange",
+    "other": "violet",
+}
 
 #label2color = {'question':'blue', 'answer':'green', 'header':'orange', 'other':'violet'}