Update app.py

app.py

@@ -18,6 +18,10 @@ import plotly.graph_objects as go
 from typing import Tuple
 import plotly.io as pio
 
+
+from wordcloud import WordCloud
+import io
+
 def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
     hex_color = hex_color.lstrip('#')
     return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))
@@ -147,19 +151,43 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
         paper_bgcolor='rgba(0,0,0,0)'
     )
     
-    return fig1, fig2
+    # Create word cloud for extended classification entities
+    wordcloud_fig = generate_wordcloud(ner_output_ext['entities'], 'tab10')
 
-@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 = generate_charts(ner_output_bin, ner_output_ext)
+    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)
     
-    return (ner_output_bin, ner_output_ext, 
-            classification_output[0], classification_output[1], classification_output[2],
-            pie_chart, bar_chart)
+    fig, ax = plt.subplots(figsize=(10, 5))
+    ax.imshow(wc, interpolation='bilinear')
+    ax.axis('off')
+    return fig
+
+@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, binary_pipeline)
+    ner_output_ext = process_ner(text, extended_pipeline)
+
+    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"]]
+
+    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"]]
+
+    fig1, fig2, wordcloud_fig = generate_charts(ner_output_bin, ner_output_ext)
+
+    # Convert the word cloud figure to an image for Gradio
+    buf = io.BytesIO()
+    wordcloud_fig.savefig(buf, format='png')
+    buf.seek(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
+
+    return bin_text, ext_text, internal_count, external_count, internal_ratio, fig1, fig2, buf
 
 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.'],
@@ -191,7 +219,8 @@ iface = gr.Interface(
         gr.Label(label="External Detail Count"),
         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.Plot(label="Binary SeqClass Entity Count Bar Chart"),
+        gr.Image(label="Word Cloud of Extended Entities")
     ],
     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.",
@@ -199,4 +228,4 @@ iface = gr.Interface(
     theme=monochrome
 )
 
-iface.launch()
+iface.launch()