Update app.py

app.py

@@ -62,28 +62,21 @@ def analyze_next_token(input_text, temperature, top_p, top_k):
     
     try:
         with torch.no_grad():
-            outputs = model(**inputs, output_attentions=True)
+            outputs = model(**inputs)
         
         last_token_logits = outputs.logits[0, -1, :]
         probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
-        
-        # Obtenir les 10 tokens les plus probables
         top_k = 10
         top_probs, top_indices = torch.topk(probabilities, top_k)
         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)
         
-        # Créer un texte explicatif
-        prob_text = "Prochains tokens les plus probables :\n\n"
-        for word, prob in prob_data.items():
-            escaped_word = word.replace("<", "&lt;").replace(">", "&gt;")
-            prob_text += f"{escaped_word}: {prob:.2%}\n"
+        prob_text = "\n".join([f"{word}: {prob:.2%}" for word, prob in prob_data.items()])
         
-        # Créer les visualisations
-        prob_plot = plot_probabilities(prob_data)
-        attention_plot = plot_attention(inputs["input_ids"][0], outputs.attentions)
+        attention_heatmap = plot_attention_alternative(inputs["input_ids"][0], last_token_logits)
         
-        return prob_text, attention_plot, prob_plot
+        return prob_text, attention_heatmap, prob_plot
     except Exception as e:
         return f"Erreur lors de l'analyse : {str(e)}", None, None
 
@@ -115,32 +108,40 @@ def plot_probabilities(prob_data):
     probs = list(prob_data.values())
     
     fig, ax = plt.subplots(figsize=(12, 6))
-    bars = ax.bar(range(len(words)), probs, color='lightgreen')
-    ax.set_title("Probabilités des tokens suivants les plus probables")
-    ax.set_xlabel("Tokens")
-    ax.set_ylabel("Probabilité")
-    
-    ax.set_xticks(range(len(words)))
-    ax.set_xticklabels(words, rotation=45, ha='right')
-    
-    for i, (bar, word) in enumerate(zip(bars, words)):
+    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(i, height, f'{word}\n{height:.2%}',
-                ha='center', va='bottom', rotation=0)
+        ax.text(bar.get_x() + bar.get_width()/2., height,
+                f'{height:.2%}',
+                ha='center', va='bottom', fontsize=10)
     
     plt.tight_layout()
     return fig
 
-def plot_attention(input_ids, attention_outputs):
+def plot_attention_alternative(input_ids, last_token_logits):
     input_tokens = tokenizer.convert_ids_to_tokens(input_ids)
+    attention_scores = torch.nn.functional.softmax(last_token_logits, dim=-1)
+    top_k = min(len(input_tokens), 10)  # Limiter à 10 tokens pour la lisibilité
+    top_attention_scores, _ = torch.topk(attention_scores, top_k)
     
-    # Prendre la moyenne des attentions sur toutes les couches et têtes
-    attention = torch.mean(torch.cat(attention_outputs), dim=(0, 1)).cpu().numpy()
+    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)
     
-    fig, ax = plt.subplots(figsize=(12, 10))
-    sns.heatmap(attention, annot=True, cmap="YlOrRd", xticklabels=input_tokens, yticklabels=input_tokens, ax=ax)
+    # Ajuster la colorbar
+    cbar = ax.collections[0].colorbar
+    cbar.set_label("Score d'attention", fontsize=12)
+    cbar.ax.tick_params(labelsize=10)
     
-    ax.set_title("Carte d'attention moyenne")
     plt.tight_layout()
     return fig
 
@@ -151,7 +152,7 @@ def reset():
     return "", 1.0, 1.0, 50, None, None, None, None
 
 with gr.Blocks() as demo:
-    gr.Markdown("# LLM & Bias")
+    gr.Markdown("# Analyse et génération de texte")
     
     with gr.Accordion("Sélection du modèle"):
         model_dropdown = gr.Dropdown(choices=models, label="Choisissez un modèle")