Update app.py

app.py

@@ -6,6 +6,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 from huggingface_hub import login
 import os
+
 login(token=os.environ["HF_TOKEN"])
 
 # Liste des modèles
@@ -50,23 +51,52 @@ def generate_text(input_text, temperature, top_p, top_k):
     # Obtenir les logits pour le dernier token généré
     last_token_logits = model(outputs.sequences[:, -1:]).logits[:, -1, :]
     
+    # Appliquer softmax pour obtenir les probabilités
+    probabilities = torch.nn.functional.softmax(last_token_logits[0], dim=-1)
+    
+    # Obtenir les top 5 tokens les plus probables
+    top_k = 5
+    top_probs, top_indices = torch.topk(probabilities, top_k)
+    top_words = [tokenizer.decode([idx.item()]) for idx in top_indices]
+    
+    # Préparer les données pour le graphique des probabilités
+    prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
+    
     # Extraire les attentions
     attentions = outputs.attentions[-1][-1].mean(dim=0).numpy()
     
-    # Visualiser l'attention
-    plt.figure(figsize=(10, 10))
-    plt.imshow(attentions, cmap='viridis')
-    plt.title("Carte d'attention")
-    attention_plot = plt.gcf()
-    plt.close()
+    # Préparer les données pour la carte d'attention
+    tokens = tokenizer.convert_ids_to_tokens(outputs.sequences[0])
+    attention_data = {
+        'attention': attentions.tolist(),
+        'tokens': tokens
+    }
     
-    # Obtenir les mots les plus probables
-    probs = torch.nn.functional.softmax(last_token_logits, dim=-1)
-    top_probs, top_indices = torch.topk(probs[0], k=5)
-    top_words = [tokenizer.decode([idx.item()]) for idx in top_indices]
+    return generated_text, attention_data, prob_data
+
+def plot_attention(attention_data):
+    attention = np.array(attention_data['attention'])
+    tokens = attention_data['tokens']
     
-    return generated_text, attention_plot, top_words
+    plt.figure(figsize=(10, 10))
+    plt.imshow(attention, cmap='viridis')
+    plt.colorbar()
+    plt.xticks(range(len(tokens)), tokens, rotation=90)
+    plt.yticks(range(len(tokens)), tokens)
+    plt.title("Carte d'attention")
+    return plt
 
+def plot_probabilities(prob_data):
+    words = list(prob_data.keys())
+    probs = list(prob_data.values())
+    
+    plt.figure(figsize=(10, 5))
+    plt.bar(words, probs)
+    plt.title("Probabilités des tokens suivants les plus probables")
+    plt.xlabel("Tokens")
+    plt.ylabel("Probabilité")
+    plt.xticks(rotation=45)
+    return plt
 
 def reset():
     return "", 1.0, 1.0, 50, None, None, None
@@ -91,15 +121,18 @@ with gr.Blocks() as demo:
     
     with gr.Row():
         attention_plot = gr.Plot(label="Visualisation de l'attention")
-        top_words = gr.JSON(label="Mots les plus probables")
+        prob_plot = gr.Plot(label="Probabilités des tokens suivants")
     
     reset_button = gr.Button("Réinitialiser")
     
     load_button.click(load_model, inputs=[model_dropdown], outputs=[load_output])
     generate_button.click(generate_text, 
                           inputs=[input_text, temperature, top_p, top_k], 
-                          outputs=[output_text, attention_plot, top_words])
+                          outputs=[output_text, attention_plot, prob_plot])
     reset_button.click(reset, 
-                       outputs=[input_text, temperature, top_p, top_k, output_text, attention_plot, top_words])
+                       outputs=[input_text, temperature, top_p, top_k, output_text, attention_plot, prob_plot])
+    
+    attention_plot.change(plot_attention, inputs=[attention_plot], outputs=[attention_plot])
+    prob_plot.change(plot_probabilities, inputs=[prob_plot], outputs=[prob_plot])
 
 demo.launch()