Update app.py

app.py

@@ -62,17 +62,18 @@ def analyze_next_token(input_text, temperature, top_p, top_k):
         probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
         top_k = 5
         top_probs, top_indices = torch.topk(probabilities, top_k)
-        top_words = [tokenizer.decode([idx.item()]) for idx in top_indices]
+        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()])
+        
+        # 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_data = torch.mean(torch.stack(outputs.attentions), dim=(0, 1)).cpu().numpy()
-            attention_plot = plot_attention(attention_data, tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]))
-        else:
-            attention_plot = None
+            attention_text = "Attention disponible"
         
-        return "\n".join([f"{word}: {prob:.4f}" for word, prob in prob_data.items()]), attention_plot, prob_plot
+        return prob_text, attention_text, prob_plot
     except Exception as e:
         return f"Erreur lors de l'analyse : {str(e)}", None, None
 
@@ -88,29 +89,19 @@ def generate_text(input_text, temperature, top_p, top_k):
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=50,
+                max_new_tokens=1,  # Génère seulement le prochain mot
                 temperature=temperature,
                 top_p=top_p,
                 top_k=top_k
             )
         
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return generated_text
+        # Ne retourne que le nouveau mot généré
+        new_word = generated_text[len(input_text):].strip()
+        return new_word
     except Exception as e:
         return f"Erreur lors de la génération : {str(e)}"
 
-def plot_attention(attention, tokens):
-    fig, ax = plt.subplots(figsize=(10, 10))
-    im = ax.imshow(attention, cmap='viridis')
-    ax.set_xticks(range(len(tokens)))
-    ax.set_yticks(range(len(tokens)))
-    ax.set_xticklabels(tokens, rotation=90)
-    ax.set_yticklabels(tokens)
-    plt.colorbar(im)
-    plt.title("Carte d'attention")
-    plt.tight_layout()
-    return fig
-
 def plot_probabilities(prob_data):
     words = list(prob_data.keys())
     probs = list(prob_data.values())
@@ -145,27 +136,26 @@ with gr.Blocks() as demo:
     
     input_text = gr.Textbox(label="Texte d'entrée", lines=3)
     analyze_button = gr.Button("Analyser le prochain token")
-    generate_button = gr.Button("Générer la suite du texte")
     
     next_token_probs = gr.Textbox(label="Probabilités du prochain token")
+    attention_info = gr.Textbox(label="Information sur l'attention")
     
-    with gr.Row():
-        attention_plot = gr.Plot(label="Visualisation de l'attention")
-        prob_plot = gr.Plot(label="Probabilités des tokens suivants")
+    prob_plot = gr.Plot(label="Probabilités des tokens suivants")
     
-    generated_text = gr.Textbox(label="Texte généré", lines=5)
+    generate_button = gr.Button("Générer le prochain mot")
+    generated_word = gr.Textbox(label="Mot généré")
     
     reset_button = gr.Button("Réinitialiser")
     
     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, attention_info, prob_plot])
     generate_button.click(generate_text, 
                           inputs=[input_text, temperature, top_p, top_k], 
-                          outputs=[generated_text])
+                          outputs=[generated_word])
     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, attention_info, prob_plot, generated_word])
 
 if __name__ == "__main__":
     demo.launch()