Update app.py

app.py

@@ -1,14 +1,12 @@
 from __future__ import annotations
 from typing import Iterable, Tuple
-
 import gradio as gr
 from gradio.themes.monochrome import Monochrome
 import spaces
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, pipeline
 import os
-import matplotlib.pyplot as plt
-import colorsys
+import plotly.graph_objects as go
 
 # Utility functions for color conversions and brightness adjustment
 def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
@@ -85,15 +83,11 @@ def process_classification(text, model1, model2, tokenizer1) -> Tuple[float, flo
 
     return prediction1, prediction2, score
 
-def generate_pie_chart(internal_count: float, external_count: float):
-    labels = 'Internal', 'External'
-    sizes = [internal_count, external_count]
-    colors = ['#ff9999','#66b3ff']
-    fig, ax = plt.subplots()
-    ax.pie(sizes, labels=labels, colors=colors, autopct='%1.1f%%', startangle=140)
-    plt.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
-    file_path = "/tmp/pie_chart.png"
-    plt.savefig(file_path)
+def generate_pie_chart(values: list, labels: list, title: str):
+    fig = go.Figure(data=[go.Pie(labels=labels, values=values, hole=.3)])
+    fig.update_layout(title_text=title)
+    file_path = "/tmp/pie_chart.html"
+    fig.write_html(file_path)
     return file_path
 
 @spaces.GPU
@@ -101,8 +95,14 @@ def all(text):
     ner_bin = process_ner(text, pipe_bin)
     ner_ext = process_ner(text, pipe_ext)
     int_count, ext_count, int_ext_ratio = process_classification(text, model1, model2, tokenizer1)
-    pie_chart_path = generate_pie_chart(int_count, ext_count)
-    return ner_bin, ner_ext, f"{round(int_count, 1)}", f"{round(ext_count, 1)}", f"{round(int_ext_ratio, 2)}", pie_chart_path
+    
+    # Create pie charts
+    pie_chart_path_int_ext = generate_pie_chart([int_count, ext_count], ['Internal', 'External'], "Internal vs External Details")
+    subclass_labels = [entity['entity'] for entity in ner_ext['entities']]
+    subclass_values = [1] * len(subclass_labels)  # Each entity is counted once; adjust as needed for actual counts
+    pie_chart_path_subclass = generate_pie_chart(subclass_values, subclass_labels, "Detail Subclasses")
+
+    return ner_bin, ner_ext, f"{round(int_count, 1)}", f"{round(ext_count, 1)}", f"{round(int_ext_ratio, 2)}", pie_chart_path_int_ext, pie_chart_path_subclass
 
 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.'],
@@ -115,14 +115,15 @@ examples = [
 # Define Gradio interface
 iface = gr.Interface(
     fn=all,
-    inputs=gr.Textbox(lines=5, label="Input Text", placeholder="Write about how your breakfast went or anything else that happend or might happen to you ..."),
+    inputs=gr.Textbox(lines=5, label="Input Text", placeholder="Write about an experience"),
     outputs=[
         gr.HighlightedText(label="NER Binary"),
         gr.HighlightedText(label="NER External"),
         gr.Textbox(label="Internal Detail"),
         gr.Textbox(label="External Detail"),
         gr.Textbox(label="Ratio Int/Ext"),
-        gr.Image(type="filepath", label="Pie Chart of Internal vs External Details"),
+        gr.HTML(label="Pie Chart of Internal vs External Details", elem_id="pie_chart_int_ext"),
+        gr.HTML(label="Pie Chart of Subclass Details", elem_id="pie_chart_subclass"),
     ],
     theme=monochrome,
     examples=examples