Update app.py

change to display tokens instead

app.py

@@ -155,16 +155,23 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
     
     return fig1, fig2, wordcloud_image
 
+
 def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str]) -> np.ndarray:
-    entity_texts = [entity['entity'] for entity in entities]
-    entity_scores = [entity['score'] for entity in entities]
-    entity_types = [entity['entity'] for entity in entities]
+    token_texts = []
+    token_scores = []
+    token_types = []
+
+    for entity in entities:
+        for token in entity['tokens']:
+            token_texts.append(token)
+            token_scores.append(entity['score'])
+            token_types.append(entity['entity'])
 
     # Create a dictionary for word cloud
-    word_freq = {text: score for text, score in zip(entity_texts, entity_scores)}
+    word_freq = {text: score for text, score in zip(token_texts, token_scores)}
 
     def color_func(word, font_size, position, orientation, random_state=None, **kwargs):
-        entity_type = next(entity['entity'] for entity in entities if entity['entity'] == word)
+        entity_type = next((t for t, w in zip(token_types, token_texts) if w == word), None)
         return color_map.get(entity_type, "#FFFFFF")
 
     wordcloud = WordCloud(width=800, height=400, background_color='black', color_func=color_func).generate_from_frequencies(word_freq)