Update app.py

app.py

@@ -66,13 +66,13 @@ def analyze_next_token(input_text, temperature, top_p, top_k):
         
         last_token_logits = outputs.logits[0, -1, :]
         probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
-        top_k = 10
+        top_k = 5
         top_probs, top_indices = torch.topk(probabilities, top_k)
-        top_words = [tokenizer.decode(idx.item()).strip() for idx in top_indices]
+        top_words = [tokenizer.decode([idx.item()]).strip() for idx in top_indices]
         prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
         prob_plot = plot_probabilities(prob_data)
         
-        prob_text = "\n".join([f"{word}: {prob:.2%}" for word, prob in prob_data.items()])
+        prob_text = "\n".join([f"{word}: {prob:.4f}" for word, prob in prob_data.items()])
         
         attention_heatmap = plot_attention_alternative(inputs["input_ids"][0], last_token_logits)
         
@@ -107,21 +107,12 @@ def plot_probabilities(prob_data):
     words = list(prob_data.keys())
     probs = list(prob_data.values())
     
-    fig, ax = plt.subplots(figsize=(12, 6))
-    bars = ax.bar(words, probs, color='skyblue')
-    ax.set_title("Probabilités des 10 tokens suivants les plus probables", fontsize=16)
-    ax.set_xlabel("Tokens", fontsize=12)
-    ax.set_ylabel("Probabilité", fontsize=12)
-    plt.xticks(rotation=45, ha='right', fontsize=10)
-    plt.yticks(fontsize=10)
-    
-    # Ajouter les pourcentages au-dessus des barres
-    for bar in bars:
-        height = bar.get_height()
-        ax.text(bar.get_x() + bar.get_width()/2., height,
-                f'{height:.2%}',
-                ha='center', va='bottom', fontsize=10)
-    
+    fig, ax = plt.subplots(figsize=(10, 5))
+    sns.barplot(x=words, y=probs, ax=ax)
+    ax.set_title("Probabilités des tokens suivants les plus probables")
+    ax.set_xlabel("Tokens")
+    ax.set_ylabel("Probabilité")
+    plt.xticks(rotation=45)
     plt.tight_layout()
     return fig
 
@@ -131,17 +122,11 @@ def plot_attention_alternative(input_ids, last_token_logits):
     top_k = min(len(input_tokens), 10)  # Limiter à 10 tokens pour la lisibilité
     top_attention_scores, _ = torch.topk(attention_scores, top_k)
     
-    fig, ax = plt.subplots(figsize=(14, 7))
-    sns.heatmap(top_attention_scores.unsqueeze(0).numpy(), annot=True, cmap="YlOrRd", cbar=True, ax=ax, fmt='.2%')
-    ax.set_xticklabels(input_tokens[-top_k:], rotation=45, ha="right", fontsize=10)
-    ax.set_yticklabels(["Attention"], rotation=0, fontsize=10)
-    ax.set_title("Scores d'attention pour les derniers tokens", fontsize=16)
-    
-    # Ajuster la colorbar
-    cbar = ax.collections[0].colorbar
-    cbar.set_label("Score d'attention", fontsize=12)
-    cbar.ax.tick_params(labelsize=10)
-    
+    fig, ax = plt.subplots(figsize=(12, 6))
+    sns.heatmap(top_attention_scores.unsqueeze(0).numpy(), annot=True, cmap="YlOrRd", cbar=False, ax=ax)
+    ax.set_xticklabels(input_tokens[-top_k:], rotation=45, ha="right")
+    ax.set_yticklabels(["Attention"], rotation=0)
+    ax.set_title("Scores d'attention pour les derniers tokens")
     plt.tight_layout()
     return fig