Update app.py

app.py

@@ -66,17 +66,23 @@ 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 = 5
+        
+        # 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)
         
-        prob_text = "\n".join([f"{word}: {prob:.4f}" for word, prob in prob_data.items()])
+        # 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"
         
-        attention_heatmap = plot_attention_alternative(inputs["input_ids"][0], last_token_logits)
+        # Créer les visualisations
+        prob_plot = plot_probabilities(prob_data)
+        importance_plot = plot_token_importance(inputs["input_ids"][0], last_token_logits)
         
-        return prob_text, attention_heatmap, prob_plot
+        return prob_text, importance_plot, prob_plot
     except Exception as e:
         return f"Erreur lors de l'analyse : {str(e)}", None, None
 
@@ -108,25 +114,51 @@ def plot_probabilities(prob_data):
     probs = list(prob_data.values())
     
     fig, ax = plt.subplots(figsize=(10, 5))
-    sns.barplot(x=words, y=probs, ax=ax)
+    bars = ax.bar(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)
+    plt.xticks(rotation=45, ha='right')
+    
+    # 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')
+    
     plt.tight_layout()
     return fig
 
-def plot_attention_alternative(input_ids, last_token_logits):
+def plot_token_importance(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)
-    
-    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")
+    
+    # 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))
+    
+    # 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')
+    
     plt.tight_layout()
     return fig
 
@@ -137,7 +169,7 @@ def reset():
     return "", 1.0, 1.0, 50, None, None, None, None
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Analyse et génération de texte")
+    gr.Markdown("# LLM & Bias ")
     
     with gr.Accordion("Sélection du modèle"):
         model_dropdown = gr.Dropdown(choices=models, label="Choisissez un modèle")
@@ -155,7 +187,7 @@ with gr.Blocks() as demo:
     next_token_probs = gr.Textbox(label="Probabilités du prochain token")
     
     with gr.Row():
-        attention_plot = gr.Plot(label="Visualisation de l'attention")
+        importance_plot = gr.Plot(label="Importance des tokens d'entrée")
         prob_plot = gr.Plot(label="Probabilités des tokens suivants")
     
     generate_button = gr.Button("Générer le prochain mot")
@@ -166,12 +198,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, attention_plot, prob_plot])
+                         outputs=[next_token_probs, importance_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, attention_plot, prob_plot, generated_text])
+                       outputs=[input_text, temperature, top_p, top_k, next_token_probs, importance_plot, prob_plot, generated_text])
 
 if __name__ == "__main__":
     demo.launch()