Update app.py

app.py

@@ -122,22 +122,21 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_pat
     mask_image = np.array(Image.open(image_path))
     mask_height, mask_width = mask_image.shape[:2]
     
-    token_texts = []
-    token_scores = []
-    token_types = []
+    entity_texts = []
+    entity_scores = []
+    entity_types = []
 
     for entity in entities:
-        for token in entity['tokens']:
-            cleaned_token = re.sub(r'^\W+', '', token)
-            token_texts.append(cleaned_token)
-            token_scores.append(entity['score'])
-            token_types.append(entity['entity'])
-            print(f"{cleaned_token} ({entity['entity']}): {entity['score']}")
+        cleaned_entity = re.sub(r'^\W+', '', ' '.join(entity['tokens']))
+        entity_texts.append(cleaned_entity)
+        entity_scores.append(np.mean([token['score'] for token in entity['tokens']]))
+        entity_types.append(entity['entity'])
+        print(f"{cleaned_entity} ({entity['entity']}): {np.mean([token['score'] for token in entity['tokens']])}")
 
-    word_freq = {text: score for text, score in zip(token_texts, token_scores)}
+    word_freq = {text: score for text, score in zip(entity_texts, entity_scores)}
 
     def color_func(word, font_size, position, orientation, random_state=None, **kwargs):
-        entity_type = next((t for t, w in zip(token_types, token_texts) if w == word), None)
+        entity_type = next((t for t, w in zip(entity_types, entity_texts) if w == word), None)
         return color_map.get(entity_type, "#FFFFFF")
 
     wordcloud = WordCloud(width=mask_width, height=mask_height, background_color='#121212', mask=mask_image, color_func=color_func).generate_from_frequencies(word_freq)