Update app.py

app.py

@@ -51,11 +51,9 @@ def process_ner(text: str, pipeline) -> dict:
     output = pipeline(text)
     entities = []
     current_entity = None
-
     for token in output:
         entity_type = token['entity'][2:]
         entity_prefix = token['entity'][:1]
-
         if current_entity is None or entity_type != current_entity['entity'] or (entity_prefix == 'B' and entity_type == current_entity['entity']):
             if current_entity is not None:
                 entities.append(current_entity)
@@ -63,98 +61,80 @@ def process_ner(text: str, pipeline) -> dict:
                 "entity": entity_type,
                 "start": token['start'],
                 "end": token['end'],
-                "score": token['score'],
+                "scores": [token['score']],
                 "tokens": [token['word']]
             }
         else:
             current_entity['end'] = token['end']
-            current_entity['score'] = max(current_entity['score'], token['score'])
+            current_entity['scores'].append(token['score'])
             current_entity['tokens'].append(token['word'])
-
     if current_entity is not None:
         entities.append(current_entity)
-
+    for entity in entities:
+        entity['average_score'] = sum(entity['scores']) / len(entity['scores'])
     return {"text": text, "entities": entities}
 
 def generate_charts(ner_output_ext: dict) -> Tuple[go.Figure, np.ndarray]:
     entities_ext = [entity['entity'] for entity in ner_output_ext['entities']]
-    
-    # Counting entities for extended classification
     entity_counts_ext = {entity: entities_ext.count(entity) for entity in set(entities_ext)}
     ext_labels = list(entity_counts_ext.keys())
     ext_sizes = list(entity_counts_ext.values())
-
     ext_color_map = {
-        "INTemothou": "#FF7F50",  # Coral
-        "INTpercept": "#FF4500",  # OrangeRed
-        "INTtime": "#FF6347",     # Tomato
-        "INTplace": "#FFD700",    # Gold
-        "INTevent": "#FFA500",    # Orange
-        "EXTsemantic": "#4682B4", # SteelBlue
-        "EXTrepetition": "#5F9EA0", # CadetBlue
-        "EXTother": "#00CED1",    # DarkTurquoise
+        "INTemothou": "#FF7F50",
+        "INTpercept": "#FF4500",
+        "INTtime": "#FF6347",
+        "INTplace": "#FFD700",
+        "INTevent": "#FFA500",
+        "EXTsemantic": "#4682B4",
+        "EXTrepetition": "#5F9EA0",
+        "EXTother": "#00CED1",
     }
-
     ext_colors = [ext_color_map.get(label, "#FFFFFF") for label in ext_labels]
-    
-    # Create pie chart for extended classification
     fig1 = go.Figure(data=[go.Pie(labels=ext_labels, values=ext_sizes, textinfo='label+percent', hole=.3, marker=dict(colors=ext_colors))])
     fig1.update_layout(
         template='plotly_dark',
         plot_bgcolor='rgba(0,0,0,0)',
         paper_bgcolor='rgba(0,0,0,0)'
     )
-    # Generate word cloud
     wordcloud_image = generate_wordcloud(ner_output_ext['entities'], ext_color_map, "dh3.png")
-    
     return fig1, wordcloud_image
 
 def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_path: str) -> np.ndarray:
     # 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]
-    
-    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']}")
-
-    word_freq = {text: score for text, score in zip(token_texts, token_scores)}
-
+        print(f"E: {entity}")
+        segment_text = ' '.join(entity['tokens'])
+        #segment_text = re.sub(r'^\W+', '', segment_text)
+        entity_texts.append(segment_text)
+        if 'average_score' in entity:
+            entity_scores.append(entity['average_score'])
+        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 = {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)
-
     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