Update app.py

app.py

@@ -32,10 +32,11 @@ def generate_text(input_text, temperature, top_p, top_k):
     global model, tokenizer
     
     inputs = tokenizer(input_text, return_tensors="pt")
+    input_ids = inputs["input_ids"]
     
     with torch.no_grad():
         outputs = model.generate(
-            **inputs,
+            input_ids,
             max_new_tokens=50,
             temperature=temperature,
             top_p=top_p,
@@ -46,9 +47,11 @@ def generate_text(input_text, temperature, top_p, top_k):
     
     generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
     
-    # Extraire les attentions et les logits
-    attentions = outputs.attentions[-1][0][-1].numpy()
-    logits = outputs.scores[-1][0]
+    # Obtenir les logits pour le dernier token généré
+    last_token_logits = model(outputs.sequences[:, -1:]).logits[:, -1, :]
+    
+    # Extraire les attentions
+    attentions = outputs.attentions[-1][-1].mean(dim=0).numpy()
     
     # Visualiser l'attention
     plt.figure(figsize=(10, 10))
@@ -58,12 +61,13 @@ def generate_text(input_text, temperature, top_p, top_k):
     plt.close()
     
     # Obtenir les mots les plus probables
-    probs = torch.nn.functional.softmax(logits, dim=-1)
-    top_probs, top_indices = torch.topk(probs, k=5)
-    top_words = [tokenizer.decode([idx]) for idx in top_indices]
+    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_plot, top_words
 
+
 def reset():
     return "", 1.0, 1.0, 50, None, None, None