Update app.py

app.py

@@ -14,6 +14,9 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification, Auto
 import os
 import colorsys
 import matplotlib.pyplot as plt
+import plotly.graph_objects as go
+from typing import Tuple
+import plotly.io as pio
 
 def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
     hex_color = hex_color.lstrip('#')
@@ -86,26 +89,32 @@ def process_classification(text: str, model1, model2, tokenizer1) -> Tuple[str,
     score = prediction1 / (prediction2 + prediction1)
 
     return f"{round(prediction1, 1)}", f"{round(prediction2, 1)}", f"{round(score, 2)}"
-    
-def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[plt.Figure, plt.Figure]:
+
+def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figure, go.Figure]:
     entities_bin = [entity['entity'] for entity in ner_output_bin['entities']]
     entities_ext = [entity['entity'] for entity in ner_output_ext['entities']]
     
-    all_entities = entities_bin + entities_ext
-    entity_counts = {entity: all_entities.count(entity) for entity in set(all_entities)}
+    # Counting entities for binary classification
+    entity_counts_bin = {entity: entities_bin.count(entity) for entity in set(entities_bin)}
+    bin_labels = list(entity_counts_bin.keys())
+    bin_sizes = list(entity_counts_bin.values())
     
-    pie_labels = list(entity_counts.keys())
-    pie_sizes = list(entity_counts.values())
-
-    fig1, ax1 = plt.subplots()
-    ax1.pie(pie_sizes, labels=pie_labels, autopct='%1.1f%%', startangle=90)
-    ax1.axis('equal')
+    # 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())
+    
+    # Create pie chart for extended classification
+    fig1 = go.Figure(data=[go.Pie(labels=ext_labels, values=ext_sizes, textinfo='label+percent', hole=.3)])
+    fig1.update_layout(title_text='Extended Sequence Classification Subclasses')
     
-    fig2, ax2 = plt.subplots()
-    ax2.bar(entity_counts.keys(), entity_counts.values())
-    ax2.set_ylabel('Count')
-    ax2.set_xlabel('Entity Type')
-    ax2.set_title('Entity Counts')
+    # Create bar chart for binary classification
+    fig2 = go.Figure(data=[go.Bar(x=bin_labels, y=bin_sizes)])
+    fig2.update_layout(
+        title='Binary Sequence Classification Classes',
+        xaxis_title='Entity Type',
+        yaxis_title='Count'
+    )
     
     return fig1, fig2