Update app.py

app.py

@@ -6,11 +6,11 @@ import numpy as np
 import html
 from transformers import AutoTokenizer, AutoModel, logging as hf_logging
 import pandas as pd
-import matplotlib # Still used for a basic empty plot if Plotly one is too complex for that
-matplotlib.use('Agg') # Matplotlib backend setting
+import matplotlib
+matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 from sklearn.decomposition import PCA
-import plotly.graph_objects as go # For interactive 3D PCA plot
+import plotly.graph_objects as go
 
 # --- Global Settings and Model Loading ---
 hf_logging.set_verbosity_error()
@@ -103,7 +103,6 @@ def plot_token_pca_3d_plotly(token_embeddings_3d, tokens, scores, title="Token E
     fig.update_layout(
         title=dict(text=title, x=0.5, font=dict(size=16)),
         scene=dict(
-            # 수정된 부분: title 속성 내에 text와 font를 포함
             xaxis=dict(title=dict(text='PCA Comp 1', font=dict(size=10)), tickfont=dict(size=9), backgroundcolor="rgba(230, 230, 230, 0.8)"),
             yaxis=dict(title=dict(text='PCA Comp 2', font=dict(size=10)), tickfont=dict(size=9), backgroundcolor="rgba(230, 230, 230, 0.8)"),
             zaxis=dict(title=dict(text='PCA Comp 3', font=dict(size=10)), tickfont=dict(size=9), backgroundcolor="rgba(230, 230, 230, 0.8)"),
@@ -134,8 +133,9 @@ def analyze_sentence_for_gradio(sentence_text, top_k_value):
         error_html = f"<p style='color:red;'>Initialization Error: {html.escape(MODEL_LOADING_ERROR_MESSAGE)}</p>"
         empty_df = pd.DataFrame(columns=['token', 'score'])
         empty_fig = create_empty_plotly_figure("Model Loading Failed")
-        # gr.Label에 대한 오류 반환값 수정
-        return error_html, [], "Model Loading Failed", {"Status":"Error", "Message":"Model Loading Failed"}, [], empty_df, empty_fig
+        # gr.Label에 대한 오류 반환값 수정 (단순 딕셔너리 또는 문자열)
+        error_label_output = {"Status": "Error", "Message": "Model Loading Failed. Check logs."}
+        return error_html, [], "Model Loading Failed", error_label_output, [], empty_df, empty_fig
 
     try:
         tokenizer, model = TOKENIZER_GLOBAL, MODEL_GLOBAL
@@ -147,7 +147,8 @@ def analyze_sentence_for_gradio(sentence_text, top_k_value):
         if input_ids.shape[1] == 0:
             empty_df = pd.DataFrame(columns=['token', 'score'])
             empty_fig = create_empty_plotly_figure("Invalid Input")
-            return "<p style='color:orange;'>Input Error: No valid tokens found.</p>", [], "Input Error", {"Status":"Error", "Message":"Invalid Input"}, [], empty_df, empty_fig
+            error_label_output = {"Status": "Error", "Message": "Invalid input, no valid tokens."}
+            return "<p style='color:orange;'>Input Error: No valid tokens found.</p>", [], "Input Error", error_label_output, [], empty_df, empty_fig
 
         input_embeds_detached = model.embeddings.word_embeddings(input_ids).clone().detach()
         input_embeds_for_grad = input_embeds_detached.clone().requires_grad_(True)
@@ -166,7 +167,8 @@ def analyze_sentence_for_gradio(sentence_text, top_k_value):
         if input_embeds_for_grad.grad is None:
             empty_df = pd.DataFrame(columns=['token', 'score'])
             empty_fig = create_empty_plotly_figure("Gradient Error")
-            return "<p style='color:red;'>Analysis Error: Gradient calculation failed.</p>", [],"Analysis Error", {"Status":"Error", "Message":"Gradient Error"}, [], empty_df, empty_fig
+            error_label_output = {"Status": "Error", "Message": "Gradient calculation failed."}
+            return "<p style='color:red;'>Analysis Error: Gradient calculation failed.</p>", [],"Analysis Error", error_label_output, [], empty_df, empty_fig
         
         grads = input_embeds_for_grad.grad.clone().detach()
         scores = (grads * input_embeds_detached).norm(dim=2).squeeze(0)
@@ -210,11 +212,12 @@ def analyze_sentence_for_gradio(sentence_text, top_k_value):
         
         barplot_df = pd.DataFrame(top_tokens_for_barplot_list) if top_tokens_for_barplot_list else pd.DataFrame(columns=['token', 'score'])
             
-        predicted_class_label_str = CLASS_LABEL_MAP.get(pred_idx, f"Unknown Index: {pred_idx}")
+        predicted_class_label_str = CLASS_LABEL_MAP.get(pred_idx, f"Unknown Index ({pred_idx})")
         
         prediction_summary_text = f"Predicted Class: {predicted_class_label_str}\nProbability: {pred_prob_val:.3f}"
-        prediction_details_for_label = {"Predicted Class": predicted_class_label_str, "Probability": f"{pred_prob_val:.3f}"}
-        
+        # 수정된 부분: gr.Label에 적합한 딕셔너리 형태 (클래스명: 확률값)
+        prediction_details_for_label = {predicted_class_label_str: float(f"{pred_prob_val:.3f}")} # 확률값을 float으로 전달
+
         pca_fig = create_empty_plotly_figure("PCA Plot N/A\n(Not enough non-special tokens for 3D)")
         non_special_token_indices = [idx for idx, token_id in enumerate(input_ids[0,:len(actual_tokens)].tolist())
                                      if token_id not in [cls_token_id, sep_token_id]]
@@ -242,7 +245,8 @@ def analyze_sentence_for_gradio(sentence_text, top_k_value):
         empty_df = pd.DataFrame(columns=['token', 'score'])
         empty_fig = create_empty_plotly_figure("Analysis Error")
         # gr.Label에 대한 오류 반환값 수정
-        return error_html, [], "Analysis Failed", {"Status":"Error", "Message": str(e)}, [], empty_df, empty_fig
+        error_label_output = {"Status": "Error", "Message": f"Analysis failed: {str(e)}"}
+        return error_html, [], "Analysis Failed", error_label_output, [], empty_df, empty_fig
 
 # --- Gradio UI Definition (Translated and Enhanced) ---
 theme = gr.themes.Monochrome(
@@ -272,7 +276,7 @@ with gr.Blocks(title="AI Sentence Analyzer XAI 🚀", theme=theme, css=".gradio-
         with gr.Column(scale=2):
             with gr.Accordion("🎯 Prediction Outcome", open=True):
                 output_prediction_summary = gr.Textbox(label="Prediction Summary", lines=2, interactive=False)
-                output_prediction_details = gr.Label(label="Detailed Prediction")
+                output_prediction_details = gr.Label(label="Prediction Details & Confidence") # 레이블 이름 변경
             with gr.Accordion("⭐ Top-K Important Tokens (Table)", open=True):
                 output_top_tokens_df = gr.DataFrame(headers=["Token", "Score"], label="Most Important Tokens", 
                                                     row_count=(1,"dynamic"), col_count=(2,"fixed"), interactive=False, wrap=True)