Update app.py

app.py

@@ -109,34 +109,39 @@ 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]
 
-    entity_texts = []
-    entity_scores = []
-    entity_types = []
+    word_details = []
 
     for entity in entities:
-        print(f"E: {entity}")
-        # Join tokens and process the text
-        segment_text = ' '.join(entity['tokens'])
-        # Replace "▁" with normal spaces
-        segment_text = segment_text.replace("▁", " ")
-        # Remove all normal spaces
-        segment_text = segment_text.replace(" ", "")
-        
-        entity_texts.append(segment_text)
-        if 'average_score' in entity:
-            entity_scores.append(entity['average_score'])
+        for token in entity['tokens']:
+            # Process each token
+            token_text = token.replace("▁", " ").strip()
+            if token_text:  # Ensure token is not empty
+                word_details.append({
+                    'text': token_text,
+                    'score': entity.get('average_score', 0.5),
+                    'entity': entity['entity']
+                })
+
+    # Calculate word frequency weighted by score
+    word_freq = {}
+    for detail in word_details:
+        if detail['text'] in word_freq:
+            word_freq[detail['text']]['score'] += detail['score']
+            word_freq[detail['text']]['count'] += 1
         else:
-            entity_scores.append(0.5)  # Example: Assigning a default score of 0.5 if 'average_score' is missing
-        entity_types.append(entity['entity'])
-        print(f"{segment_text} ({entity['entity']}): {entity.get('average_score', 0.5)}")  # Print or log the score, using a default if missing
+            word_freq[detail['text']] = {'score': detail['score'], 'count': 1, 'entity': detail['entity']}
 
-    word_freq = {text: score for text, score in zip(entity_texts, entity_scores)}
+    # Average the scores and prepare final frequency dictionary
+    final_word_freq = {word: details['score'] / details['count'] for word, details in word_freq.items()}
+
+    # Prepare entity type mapping for color function
+    word_to_entity = {word: details['entity'] for word, details in word_freq.items()}
 
     def color_func(word, font_size, position, orientation, random_state=None, **kwargs):
-        entity_type = next((t for t, w in zip(entity_types, entity_texts) if w == word), None)
+        entity_type = word_to_entity.get(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)
+    wordcloud = WordCloud(width=mask_width, height=mask_height, background_color='#121212', mask=mask_image, color_func=color_func).generate_from_frequencies(final_word_freq)
 
     plt.figure(figsize=(mask_width/100, mask_height/100))
     plt.imshow(wordcloud, interpolation='bilinear')