Update app.py

app.py

@@ -105,15 +105,23 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_pat
     image_path = os.path.join(base_path, file_path)
     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 = []
+
     for entity in entities:
         print(f"E: {entity}")
+        # Join tokens and process the text
         segment_text = ' '.join(entity['tokens'])
-        #segment_text = re.sub(r'^\W+', '', segment_text)
+        # 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'])
@@ -121,20 +129,26 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_pat
             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 = {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(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)
+
     plt.figure(figsize=(mask_width/100, mask_height/100))
     plt.imshow(wordcloud, interpolation='bilinear')
     plt.axis('off')
     plt.tight_layout(pad=0)
+
     plt_image = plt.gcf()
     plt_image.canvas.draw()
     image_array = np.frombuffer(plt_image.canvas.tostring_rgb(), dtype=np.uint8)
     image_array = image_array.reshape(plt_image.canvas.get_width_height()[::-1] + (3,))
     plt.close()
+
     return image_array
 
 @spaces.GPU