Update app.py

app.py

@@ -86,26 +86,24 @@ 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]:
-    entities_bin = [entity['entity'] for entity in ner_output_bin['entities']]
+
+def generate_charts(ner_output_bin: dict, ner_output_ext: dict, internal_count: float, external_count: float, score: float) -> Tuple[plt.Figure, plt.Figure]:
     entities_ext = [entity['entity'] for entity in ner_output_ext['entities']]
+    entity_counts_ext = {entity: entities_ext.count(entity) for entity in set(entities_ext)}
     
-    all_entities = entities_bin + entities_ext
-    entity_counts = {entity: all_entities.count(entity) for entity in set(all_entities)}
-    
-    pie_labels = list(entity_counts.keys())
-    pie_sizes = list(entity_counts.values())
+    pie_labels = list(entity_counts_ext.keys())
+    pie_sizes = list(entity_counts_ext.values())
 
     fig1, ax1 = plt.subplots()
     ax1.pie(pie_sizes, labels=pie_labels, autopct='%1.1f%%', startangle=90)
     ax1.axis('equal')
     
     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')
+    bars = ['Internal Detail Count', 'External Detail Count', 'Binary Classification Score']
+    values = [internal_count, external_count, float(score)]
+    ax2.bar(bars, values)
+    ax2.set_ylabel('Count/Score')
+    ax2.set_title('Internal vs External Details and Classification Score')
     
     return fig1, fig2
 
@@ -113,12 +111,12 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[plt.Fig
 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)
+    internal_count, external_count, score = process_classification(text, model1, model2, tokenizer1)
     
-    pie_chart, bar_chart = generate_charts(ner_output_bin, ner_output_ext)
+    pie_chart, bar_chart = generate_charts(ner_output_bin, ner_output_ext, float(internal_count), float(external_count), score)
     
     return (ner_output_bin, ner_output_ext, 
-            classification_output[0], classification_output[1], classification_output[2],
+            internal_count, external_count, score,
             pie_chart, bar_chart)
 
 examples = [
@@ -150,8 +148,8 @@ iface = gr.Interface(
         gr.Label(label="Internal Detail Count"),
         gr.Label(label="External Detail Count"),
         gr.Label(label="Approximated Internal Detail Ratio"),
-        gr.Plot(label="Entity Distribution Pie Chart"),
-        gr.Plot(label="Entity Count Bar Chart")
+        gr.Plot(label="Extended Sequence Classification Pie Chart"),
+        gr.Plot(label="Internal vs External Details and Classification Score Bar Chart")
     ],
     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.",