Update app.py

app.py

@@ -5,6 +5,8 @@ from huggingface_hub import login
 import os
 import matplotlib.pyplot as plt
 import seaborn as sns
+import numpy as np
+import time
 
 # Authentification
 login(token=os.environ["HF_TOKEN"])
@@ -35,32 +37,33 @@ def load_model(model_name, progress=gr.Progress()):
     try:
         progress(0, desc="Chargement du tokenizer")
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        progress(0.3, desc="Tokenizer chargé")
-        
-        progress(0.3, desc="Chargement du modèle")
+        progress(0.5, desc="Chargement du modèle")
         model = AutoModelForCausalLM.from_pretrained(
             model_name, 
-            torch_dtype=torch.bfloat16, 
-            device_map="auto", 
+            torch_dtype=torch.float32,
+            device_map="cpu",
             attn_implementation="eager"
         )
-        progress(0.9, desc="Modèle chargé")
-        
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
-        
-        progress(1.0, desc="Chargement terminé")
+        progress(1.0, desc="Modèle chargé")
         return f"Modèle {model_name} chargé avec succès."
     except Exception as e:
         return f"Erreur lors du chargement du modèle : {str(e)}"
 
+def ensure_token_display(token):
+    """Assure que le token est affiché correctement."""
+    if token.isdigit() or (token.startswith('-') and token[1:].isdigit()):
+        return tokenizer.decode([int(token)])
+    return token
+
 def analyze_next_token(input_text, temperature, top_p, top_k):
     global model, tokenizer
     
     if model is None or tokenizer is None:
         return "Veuillez d'abord charger un modèle.", None, None
 
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     
     try:
         with torch.no_grad():
@@ -68,20 +71,20 @@ 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
+        
+        top_k = 10
         top_probs, top_indices = torch.topk(probabilities, top_k)
-        top_words = [tokenizer.decode([idx.item()]).strip() for idx in top_indices]
+        top_words = [ensure_token_display(tokenizer.decode([idx.item()])) 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()])
+        prob_text = "Prochains tokens les plus probables :\n\n"
+        for word, prob in prob_data.items():
+            prob_text += f"{word}: {prob:.2%}\n"
         
-        # Simplification de l'affichage de l'attention
-        attention_text = "Attention non disponible pour ce modèle"
-        if hasattr(outputs, 'attentions') and outputs.attentions is not None:
-            attention_text = "Attention disponible"
+        prob_plot = plot_probabilities(prob_data)
+        attention_plot = plot_attention(inputs["input_ids"][0], last_token_logits)
         
-        return prob_text, attention_text, prob_plot
+        return prob_text, attention_plot, prob_plot
     except Exception as e:
         return f"Erreur lors de l'analyse : {str(e)}", None, None
 
@@ -91,20 +94,20 @@ def generate_text(input_text, temperature, top_p, top_k):
     if model is None or tokenizer is None:
         return "Veuillez d'abord charger un modèle."
 
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     
     try:
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=1,  # Génère seulement le prochain mot
+                max_new_tokens=1,
                 temperature=temperature,
                 top_p=top_p,
                 top_k=top_k
             )
         
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return generated_text  # Retourne l'input + le nouveau mot généré
+        return generated_text
     except Exception as e:
         return f"Erreur lors de la génération : {str(e)}"
 
@@ -112,12 +115,39 @@ def plot_probabilities(prob_data):
     words = list(prob_data.keys())
     probs = list(prob_data.values())
     
-    fig, ax = plt.subplots(figsize=(10, 5))
-    sns.barplot(x=words, y=probs, ax=ax)
+    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)
+    
+    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(i, height, f'{height:.2%}',
+                ha='center', va='bottom', rotation=0)
+    
+    plt.tight_layout()
+    return fig
+
+def plot_attention(input_ids, last_token_logits):
+    input_tokens = [ensure_token_display(tokenizer.decode([id])) for id in input_ids]
+    attention_scores = torch.nn.functional.softmax(last_token_logits, dim=-1)
+    top_k = min(len(input_tokens), 10)
+    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)
+    
+    cbar = ax.collections[0].colorbar
+    cbar.set_label("Score d'attention", fontsize=12)
+    cbar.ax.tick_params(labelsize=10)
+    
     plt.tight_layout()
     return fig
 
@@ -144,9 +174,10 @@ with gr.Blocks() as demo:
     analyze_button = gr.Button("Analyser le prochain token")
     
     next_token_probs = gr.Textbox(label="Probabilités du prochain token")
-    attention_info = gr.Textbox(label="Information sur l'attention")
     
-    prob_plot = gr.Plot(label="Probabilités des tokens suivants")
+    with gr.Row():
+        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")
     generated_text = gr.Textbox(label="Texte généré")
@@ -156,12 +187,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_info, 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, attention_info, 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()