Update app.py

app.py

@@ -122,6 +122,7 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_pat
         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]
     
     token_texts = []
     token_scores = []
@@ -141,9 +142,9 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_pat
         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='#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(word_freq)
 
-    plt.figure(figsize=(10, 5))
+    plt.figure(figsize=(mask_width/100, mask_height/100))
     plt.imshow(wordcloud, interpolation='bilinear')
     plt.axis('off')
     plt.tight_layout(pad=0)