Update app.py

app.py

@@ -62,7 +62,7 @@ def analyze_next_token(input_text, temperature, top_p, top_k):
     
     try:
         with torch.no_grad():
-            outputs = model(**inputs)
+            outputs = model(**inputs, output_attentions=True)
         
         last_token_logits = outputs.logits[0, -1, :]
         probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
@@ -76,13 +76,14 @@ def analyze_next_token(input_text, temperature, top_p, top_k):
         # Créer un texte explicatif
         prob_text = "Prochains tokens les plus probables :\n\n"
         for word, prob in prob_data.items():
-            prob_text += f"{word}: {prob:.2%}\n"
+            escaped_word = word.replace("<", "&lt;").replace(">", "&gt;")
+            prob_text += f"{escaped_word}: {prob:.2%}\n"
         
         # Créer les visualisations
         prob_plot = plot_probabilities(prob_data)
-        importance_plot = plot_token_importance(inputs["input_ids"][0], last_token_logits)
+        attention_plot = plot_attention(inputs["input_ids"][0], outputs.attentions)
         
-        return prob_text, importance_plot, prob_plot
+        return prob_text, attention_plot, prob_plot
     except Exception as e:
         return f"Erreur lors de l'analyse : {str(e)}", None, None
 
@@ -113,52 +114,33 @@ def plot_probabilities(prob_data):
     words = list(prob_data.keys())
     probs = list(prob_data.values())
     
-    fig, ax = plt.subplots(figsize=(10, 5))
-    bars = ax.bar(words, probs, color='lightgreen')
+    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é")
-    plt.xticks(rotation=45, ha='right')
     
-    # Ajouter les valeurs sur les barres
-    for bar in bars:
+    ax.set_xticks(range(len(words)))
+    ax.set_xticklabels(words, rotation=45, ha='right')
+    
+    for i, (bar, word) in enumerate(zip(bars, words)):
         height = bar.get_height()
-        ax.text(bar.get_x() + bar.get_width()/2., height,
-                f'{height:.2%}',
-                ha='center', va='bottom')
+        ax.text(i, height, f'{word}\n{height:.2%}',
+                ha='center', va='bottom', rotation=0)
     
     plt.tight_layout()
     return fig
 
-def plot_token_importance(input_ids, last_token_logits):
+def plot_attention(input_ids, attention_outputs):
     input_tokens = tokenizer.convert_ids_to_tokens(input_ids)
     
-    # Calculer l'importance de chaque token
-    importances = torch.abs(last_token_logits).sum() / len(last_token_logits)
-    importances = importances.repeat(len(input_tokens))
-    
-    # Normaliser les importances
-    importances = importances / importances.sum()
-    
-    # Créer la figure
-    fig, ax = plt.subplots(figsize=(12, 3))
+    # Prendre la moyenne des attentions sur toutes les couches et têtes
+    attention = torch.mean(torch.cat(attention_outputs), dim=(0, 1)).cpu().numpy()
     
-    # Créer un graphique à barres
-    bars = ax.bar(range(len(input_tokens)), importances, color='skyblue')
-    
-    # Ajouter les labels et le titre
-    ax.set_xticks(range(len(input_tokens)))
-    ax.set_xticklabels(input_tokens, rotation=45, ha='right')
-    ax.set_ylabel('Importance relative')
-    ax.set_title('Importance des tokens d\'entrée pour la prédiction')
-    
-    # Ajouter les valeurs sur les 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')
+    fig, ax = plt.subplots(figsize=(12, 10))
+    sns.heatmap(attention, annot=True, cmap="YlOrRd", xticklabels=input_tokens, yticklabels=input_tokens, ax=ax)
     
+    ax.set_title("Carte d'attention moyenne")
     plt.tight_layout()
     return fig
 
@@ -169,7 +151,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")
@@ -187,7 +169,7 @@ with gr.Blocks() as demo:
     next_token_probs = gr.Textbox(label="Probabilités du prochain token")
     
     with gr.Row():
-        importance_plot = gr.Plot(label="Importance des tokens d'entrée")
+        attention_plot = gr.Plot(label="Visualisation de l'attention")
         prob_plot = gr.Plot(label="Probabilités des tokens suivants")
     
     generate_button = gr.Button("Générer le prochain mot")
@@ -198,12 +180,12 @@ with gr.Blocks() as demo:
     load_button.click(load_model, inputs=[model_dropdown], outputs=[load_output])
     analyze_button.click(analyze_next_token, 
                          inputs=[input_text, temperature, top_p, top_k], 
-                         outputs=[next_token_probs, importance_plot, prob_plot])
+                         outputs=[next_token_probs, attention_plot, prob_plot])
     generate_button.click(generate_text, 
                           inputs=[input_text, temperature, top_p, top_k], 
                           outputs=[generated_text])
     reset_button.click(reset, 
-                       outputs=[input_text, temperature, top_p, top_k, next_token_probs, importance_plot, prob_plot, generated_text])
+                       outputs=[input_text, temperature, top_p, top_k, next_token_probs, attention_plot, prob_plot, generated_text])
 
 if __name__ == "__main__":
     demo.launch()