Update app.py

app.py

@@ -121,8 +121,12 @@ def compare_models(text, top_k=5):
         }
 
     # Create a properly structured dataframe
-    # Use display tokens from the first model for the Token column
-    df_data = {"Token": results["ERNIE-4.5-PT"]["display_tokens"]}
+    # The issue is that display_tokens has one more element than tokens
+    # We need to align them properly
+    
+    # For the Token column, we'll use the display_tokens but only for the tokens that have predictions
+    # The first token in display_tokens doesn't have a prediction (since it's the first token)
+    df_data = {"Token": results["ERNIE-4.5-PT"]["display_tokens"][1:]}
     
     # Add columns for each model
     for model_name in ["ERNIE-4.5-PT", "ERNIE-4.5-Base-PT"]:
@@ -141,7 +145,7 @@ def compare_models(text, top_k=5):
     # Add bars for both models
     fig_confidence.add_trace(go.Bar(
         name='ERNIE-4.5-PT',
-        x=results["ERNIE-4.5-PT"]["display_tokens"],
+        x=results["ERNIE-4.5-PT"]["display_tokens"][1:],
         y=results["ERNIE-4.5-PT"]["confidence_values"],
         marker_color='royalblue',
         text=[format_prob(x) for x in results["ERNIE-4.5-PT"]["confidence_values"]],
@@ -151,7 +155,7 @@ def compare_models(text, top_k=5):
     
     fig_confidence.add_trace(go.Bar(
         name='ERNIE-4.5-Base-PT',
-        x=results["ERNIE-4.5-Base-PT"]["display_tokens"],
+        x=results["ERNIE-4.5-Base-PT"]["display_tokens"][1:],
         y=results["ERNIE-4.5-Base-PT"]["confidence_values"],
         marker_color='lightseagreen',
         text=[format_prob(x) for x in results["ERNIE-4.5-Base-PT"]["confidence_values"]],
@@ -184,23 +188,23 @@ def compare_models(text, top_k=5):
     
     fig_logprob.add_trace(go.Scatter(
         name='ERNIE-4.5-PT',
-        x=results["ERNIE-4.5-PT"]["display_tokens"],
+        x=results["ERNIE-4.5-PT"]["display_tokens"][1:],
         y=pt_logprobs,
         mode='lines+markers',
         line=dict(color='royalblue', width=3),
         marker=dict(size=8),
-        text=[f"LogProb: {x}<br>Token: {t}" for x, t in zip(pt_logprobs, results["ERNIE-4.5-PT"]["display_tokens"])],
+        text=[f"LogProb: {x}<br>Token: {t}" for x, t in zip(pt_logprobs, results["ERNIE-4.5-PT"]["display_tokens"][1:])],
         hoverinfo='text'
     ))
     
     fig_logprob.add_trace(go.Scatter(
         name='ERNIE-4.5-Base-PT',
-        x=results["ERNIE-4.5-Base-PT"]["display_tokens"],
+        x=results["ERNIE-4.5-Base-PT"]["display_tokens"][1:],
         y=base_logprobs,
         mode='lines+markers',
         line=dict(color='lightseagreen', width=3),
         marker=dict(size=8),
-        text=[f"LogProb: {x}<br>Token: {t}" for x, t in zip(base_logprobs, results["ERNIE-4.5-Base-PT"]["display_tokens"])],
+        text=[f"LogProb: {x}<br>Token: {t}" for x, t in zip(base_logprobs, results["ERNIE-4.5-Base-PT"]["display_tokens"][1:])],
         hoverinfo='text'
     ))
     
@@ -290,7 +294,8 @@ def compare_models(text, top_k=5):
         f"- **LogProb**: How confident the model is in predicting each token (closer to 0 is better)\n"
         f"- **Confidence**: Probability percentage for each token prediction\n"
         f"- **Level**: Visual indicator of confidence (🟢🟡🔴)\n"
-        f"- **Top-k**: What other tokens the model considered likely"
+        f"- **Top-k**: What other tokens the model considered likely\n\n"
+        f"**Note**: The first token of your input doesn't have a prediction because there's no preceding context to predict from."
     )
 
     return comparison_df, summary, fig_confidence, fig_logprob, fig_summary
@@ -412,6 +417,8 @@ with gr.Blocks(css=css, title="ERNIE Model Comparison Tool") as demo:
         **Model Differences**:
         - **ERNIE-4.5-PT**: Instruction-tuned model, better at following complex instructions
         - **ERNIE-4.5-Base-PT**: Base model, better at general language patterns
+        
+        **Note**: The first token of your input doesn't have a prediction because there's no preceding context to predict from.
         """
     )