Update app.py

app.py

@@ -79,7 +79,7 @@ def process_ner(text: str, pipeline) -> dict:
     if current_entity is not None:
         entities.append(current_entity)
 
-    return {"text": text, "entities": entities}
+    return {"entities": entities}
 
 def process_classification(text: str, model1, model2, tokenizer1) -> Tuple[str, str, str]:
     inputs1 = tokenizer1(text, max_length=512, return_tensors='pt', truncation=True, padding=True)
@@ -97,7 +97,7 @@ def process_classification(text: str, model1, model2, tokenizer1) -> Tuple[str,
 import plotly.graph_objects as go
 from typing import Tuple
 
-def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figure, go.Figure]:
+def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figure, go.Figure, np.ndarray]:
     entities_bin = [entity['entity'] for entity in ner_output_bin['entities']]
     entities_ext = [entity['entity'] for entity in ner_output_ext['entities']]
     
@@ -151,47 +151,51 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
         paper_bgcolor='rgba(0,0,0,0)'
     )
     
-    # Create word cloud for extended classification entities
-    wordcloud_fig = generate_wordcloud(ner_output_ext['entities'], 'tab10')
-
-    return fig1, fig2, wordcloud_fig
-
-def generate_wordcloud(entities: list, color_map: dict) -> plt.Figure:
-    word_freq = {entity['entity']: entity['score'] for entity in entities}
-    wc = WordCloud(width=800, height=400, background_color='black', colormap=color_map).generate_from_frequencies(word_freq)
+    # Generate word cloud
+    wordcloud_image = generate_wordcloud(ner_output_ext['entities'], ext_color_map)
     
-    fig, ax = plt.subplots(figsize=(10, 5))
-    ax.imshow(wc, interpolation='bilinear')
-    ax.axis('off')
-    return fig
+    return fig1, fig2, wordcloud_image
 
-@spaces.GPU
-def all(text: str) -> Tuple[gr.HighlightedText, gr.HighlightedText, int, int, float, go.Figure, go.Figure, plt.Figure]:
-    ner_output_bin = process_ner(text, pipe_bin)
-    ner_output_ext = process_ner(text, pipe_ext)
+def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str]) -> np.ndarray:
+    entity_texts = [entity['entity'] for entity in entities]
+    entity_scores = [entity['score'] for entity in entities]
+    entity_types = [entity['entity'] for entity in entities]
 
-    binary_entities = [{"entity": ent["entity"], "score": ent["score"]} for ent in ner_output_bin["entities"]]
-    extended_entities = [{"entity": ent["entity"], "score": ent["score"]} for ent in ner_output_ext["entities"]]
+    # Create a dictionary for word cloud
+    word_freq = {text: score for text, score in zip(entity_texts, entity_scores)}
 
-    bin_text = [{"text": text[ent["start"]:ent["end"]], "entity": ent["entity"]} for ent in ner_output_bin["entities"]]
-    ext_text = [{"text": text[ent["start"]:ent["end"]], "entity": ent["entity"]} for ent in ner_output_ext["entities"]]
+    def color_func(word, font_size, position, orientation, random_state=None, **kwargs):
+        entity_type = next(entity['entity'] for entity in entities if entity['entity'] == word)
+        return color_map.get(entity_type, "#FFFFFF")
 
-    fig1, fig2, wordcloud_fig = generate_charts(ner_output_bin, ner_output_ext)
+    wordcloud = WordCloud(width=800, height=400, background_color='black', color_func=color_func).generate_from_frequencies(word_freq)
 
-    # Convert the word cloud figure to an image for Gradio
-    buf = io.BytesIO()
-    wordcloud_fig.savefig(buf, format='png')
-    buf.seek(0)
+    # Convert to image array
+    plt.figure(figsize=(10, 5))
+    plt.imshow(wordcloud, interpolation='bilinear')
+    plt.axis('off')
+    plt.tight_layout(pad=0)
 
-    internal_count = sum(1 for ent in ner_output_bin["entities"] if ent["entity"] == "Internal")
-    external_count = sum(1 for ent in ner_output_bin["entities"] if ent["entity"] == "External")
-    internal_ratio = internal_count / (internal_count + external_count) if (internal_count + external_count) > 0 else 0
+    # Convert plt to numpy array
+    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 bin_text, ext_text, internal_count, external_count, internal_ratio, fig1, fig2, buf
+    return image_array
 
-examples = [
-    ['Bevor ich meinen Hund kaufte bin ich immer alleine durch den Park gelaufen. Gestern war ich aber mit dem Hund losgelaufen. Das Wetter war sehr schön, nicht wie sonst im Winter. Ich weiß nicht genau. Mir fällt sonst nichts dazu ein. Wir trafen auf mehrere Spaziergänger. Ein Mann mit seinem Kind. Das Kind hat ein Eis gegessen.'],
-]
+@spaces.GPU
+def all(text: str):
+    ner_output_bin = process_ner(text, pipe_bin)
+    ner_output_ext = process_ner(text, pipe_ext)
+    classification_output = process_classification(text, model1, model2, tokenizer1)
+    
+    pie_chart, bar_chart, wordcloud_image = generate_charts(ner_output_bin, ner_output_ext)
+    
+    return (ner_output_bin, ner_output_ext, 
+            classification_output[0], classification_output[1], classification_output[2],
+            pie_chart, bar_chart, wordcloud_image)
 
 iface = gr.Interface(
     fn=all,
@@ -209,7 +213,6 @@ iface = gr.Interface(
                                "INTtime": "#FF6347",     # Tomato
                                "INTplace": "#FFD700",    # Gold
                                "INTevent": "#FFA500",    # Orange
-                        
                                "EXTsemantic": "#4682B4", # SteelBlue
                                "EXTrepetition": "#5F9EA0", # CadetBlue
                                "EXTother": "#00CED1",    # DarkTurquoise
@@ -220,7 +223,7 @@ iface = gr.Interface(
         gr.Label(label="Approximated Internal Detail Ratio"),
         gr.Plot(label="Extended SeqClass Entity Distribution Pie Chart"),
         gr.Plot(label="Binary SeqClass Entity Count Bar Chart"),
-        gr.Image(label="Word Cloud of Extended Entities")
+        gr.Image(label="Entity Word Cloud")
     ],
     title="Scoring Demo",
     description="Autobiographical Memory Analysis: This demo combines two text - and two sequence classification models to showcase our automated Autobiographical Interview scoring method. Submit a narrative to see the results.",