Update app.py

app.py

@@ -111,35 +111,45 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_pat
     # Construct the absolute path
     base_path = os.path.dirname(os.path.abspath(__file__))
     image_path = os.path.join(base_path, file_path)
-    
-    # Debugging statement to print the image path
-    print(f"Image path: {image_path}")
-
-    # Check if the file exists
     if not os.path.exists(image_path):
         raise FileNotFoundError(f"Mask image file not found: {image_path}")
 
     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:
-        cleaned_entity = re.sub(r'^\W+', '', ' '.join(entity['tokens']))
-        entity_texts.append(cleaned_entity)
-        entity_scores.append(np.mean(entity['scores']))
-        entity_types.append(entity['entity'])
-        print(f"{cleaned_entity} ({entity['entity']}): {np.mean(entity['scores'])}")
+        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:
+            word_freq[detail['text']] = {'score': detail['score'], 'count': 1, 'entity': detail['entity']}
+
+    # Average the scores and prepare final frequency dictionary
+    final_word_freq = {word: details['score'] / details['count'] for word, details in word_freq.items()}
 
-    word_freq = {text: score for text, score in zip(entity_texts, entity_scores)}
+    # 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')