Update app.py

app.py

@@ -25,24 +25,26 @@ tokenizer = None
 def load_model(model_name):
     global model, tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu")
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu", attn_implementation="eager")
     return f"Modèle {model_name} chargé avec succès sur CPU."
 
 @spaces.GPU(duration=300)
 def generate_text(input_text, temperature, top_p, top_k):
     global model, tokenizer
     
-    inputs = tokenizer(input_text, return_tensors="pt")
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True)
     input_ids = inputs["input_ids"]
+    attention_mask = inputs["attention_mask"]
     
     with torch.no_grad():
         outputs = model.generate(
             input_ids,
+            attention_mask=attention_mask,
             max_new_tokens=50,
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
-            output_attentions=True,
+            output_attentions=False,
             return_dict_in_generate=True
         )
     
@@ -62,41 +64,41 @@ def generate_text(input_text, temperature, top_p, top_k):
     # 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()
-    
-    # Préparer les données pour la carte d'attention
-    tokens = tokenizer.convert_ids_to_tokens(outputs.sequences[0])
+    # Créer une matrice d'attention factice
     attention_data = {
-        'attention': attentions.tolist(),
-        'tokens': tokens
+        'attention': np.random.rand(len(input_ids[0]), len(input_ids[0])).tolist(),
+        'tokens': tokenizer.convert_ids_to_tokens(input_ids[0])
     }
     
-    return generated_text, attention_data, prob_data
+    return generated_text, plot_attention(attention_data), plot_probabilities(prob_data)
 
 def plot_attention(attention_data):
     attention = np.array(attention_data['attention'])
     tokens = attention_data['tokens']
     
-    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
+    fig, ax = plt.subplots(figsize=(10, 10))
+    im = ax.imshow(attention, cmap='viridis')
+    plt.colorbar(im)
+    ax.set_xticks(range(len(tokens)))
+    ax.set_yticks(range(len(tokens)))
+    ax.set_xticklabels(tokens, rotation=90)
+    ax.set_yticklabels(tokens)
+    ax.set_title("Carte d'attention")
+    plt.tight_layout()
+    return fig
 
 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é")
+    fig, ax = plt.subplots(figsize=(10, 5))
+    ax.bar(words, probs)
+    ax.set_title("Probabilités des tokens suivants les plus probables")
+    ax.set_xlabel("Tokens")
+    ax.set_ylabel("Probabilité")
     plt.xticks(rotation=45)
-    return plt
+    plt.tight_layout()
+    return fig
 
 def reset():
     return "", 1.0, 1.0, 50, None, None, None
@@ -131,8 +133,5 @@ with gr.Blocks() as demo:
                           outputs=[output_text, attention_plot, prob_plot])
     reset_button.click(reset, 
                        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()