Update app.py

app.py

@@ -71,11 +71,13 @@ def process_ner(text: str, pipeline) -> dict:
                 "entity": entity_type,
                 "start": token['start'],
                 "end": token['end'],
-                "score": token['score']
+                "score": token['score'],
+                "tokens": [token['word']]
             }
         else:
             current_entity['end'] = token['end']
             current_entity['score'] = max(current_entity['score'], token['score'])
+            current_entity['tokens'].append(token['word'])
 
     if current_entity is not None:
         entities.append(current_entity)
@@ -117,23 +119,16 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
     }
 
     ext_color_map = {
-        "INTemothou": "#FF7F50",  # Coral
-        "INTpercept": "#FF4500",  # OrangeRed
-        "INTtime": "#FF6347",     # Tomato
-        "INTplace": "#FFD700",    # Gold
-        "INTevent": "#FFA500",    # Orange
-        "EXTsemantic": "#4682B4", # SteelBlue
-        "EXTrepetition": "#5F9EA0", # CadetBlue
-        "EXTother": "#00CED1",    # DarkTurquoise
+        "PER": "#FF7F50",  # Coral
+        "LOC": "#4682B4",  # SteelBlue
     }
 
-    bin_colors = [bin_color_map[label] for label in bin_labels]
-    ext_colors = [ext_color_map[label] for label in ext_labels]
+    bin_colors = [bin_color_map.get(label, "#FFFFFF") for label in bin_labels]
+    ext_colors = [ext_color_map.get(label, "#FFFFFF") for label in ext_labels]
     
     # Create pie chart for extended classification
     fig1 = go.Figure(data=[go.Pie(labels=ext_labels, values=ext_sizes, textinfo='label+percent', hole=.3, marker=dict(colors=ext_colors))])
     fig1.update_layout(
-        #title_text='Extended Sequence Classification Subclasses',
         template='plotly_dark',
         plot_bgcolor='rgba(0,0,0,0)',
         paper_bgcolor='rgba(0,0,0,0)'
@@ -142,7 +137,6 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
     # Create bar chart for binary classification
     fig2 = go.Figure(data=[go.Bar(x=bin_labels, y=bin_sizes, marker=dict(color=bin_colors))])
     fig2.update_layout(
-        #title='Binary Sequence Classification Classes',
         xaxis_title='Entity Type',
         yaxis_title='Count',
         template='plotly_dark',