Update app.py

app.py

@@ -34,7 +34,12 @@ def load_model(model_name):
     global model, tokenizer
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            torch_dtype=torch.bfloat16, 
+            device_map="auto", 
+            attn_implementation="eager"
+        )
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         return f"Modèle {model_name} chargé avec succès."
@@ -64,33 +69,23 @@ def generate_text(input_text, temperature, top_p, top_k):
         
         generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
         
-        # Obtenir les logits pour le dernier token généré
-        if outputs.scores:
+        if hasattr(outputs, 'scores') and outputs.scores:
             last_token_logits = outputs.scores[-1][0]
-            
-            # Appliquer softmax pour obtenir les probabilités
             probabilities = torch.nn.functional.softmax(last_token_logits, 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 (moyenne sur toutes les couches et têtes d'attention)
-            if outputs.attentions:
-                attentions = torch.mean(torch.stack(outputs.attentions), dim=(0, 1)).cpu().numpy()
-                attention_plot = plot_attention(attentions, tokenizer.convert_ids_to_tokens(inputs["input_ids"][0]))
-            else:
-                attention_plot = None
-            
             prob_plot = plot_probabilities(prob_data)
         else:
-            attention_plot = None
             prob_plot = None
         
+        if hasattr(outputs, 'attentions') and outputs.attentions:
+            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
+        
         return generated_text, attention_plot, prob_plot
     except Exception as e:
         return f"Erreur lors de la génération : {str(e)}", None, None
@@ -139,10 +134,10 @@ with gr.Blocks() as demo:
         top_p = gr.Slider(0.1, 1.0, value=1.0, label="Top-p")
         top_k = gr.Slider(1, 100, value=50, step=1, label="Top-k")
     
-    input_text = gr.Textbox(label="Texte d'entrée")
+    input_text = gr.Textbox(label="Texte d'entrée", lines=3)
     generate_button = gr.Button("Générer")
     
-    output_text = gr.Textbox(label="Texte généré")
+    output_text = gr.Textbox(label="Texte généré", lines=5)
     
     with gr.Row():
         attention_plot = gr.Plot(label="Visualisation de l'attention")
@@ -157,4 +152,5 @@ with gr.Blocks() as demo:
     reset_button.click(reset, 
                        outputs=[input_text, temperature, top_p, top_k, output_text, attention_plot, prob_plot])
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()