Update app.py

app.py

@@ -6,9 +6,10 @@ import numpy as np
 from huggingface_hub import login
 import os
 
+# Authentification Hugging Face avec ton token d'accès
 login(token=os.environ["HF_TOKEN"])
 
-# Liste des modèles
+# Liste des modèles disponibles
 models = [
     "meta-llama/Llama-2-13b", "meta-llama/Llama-2-7b", "meta-llama/Llama-2-70b",
     "meta-llama/Meta-Llama-3-8B", "meta-llama/Llama-3.2-3B", "meta-llama/Llama-3.1-8B",
@@ -17,7 +18,7 @@ models = [
     "croissantllm/CroissantLLMBase"
 ]
 
-# Variables globales pour stocker le modèle et le tokenizer
+# Variables pour le modèle et le tokenizer
 model = None
 tokenizer = None
 
@@ -26,14 +27,14 @@ def load_model(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
     
-    # Définir le token de padding s'il n'existe pas
+    # Assurer que le token de padding est défini si nécessaire
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
         model.config.pad_token_id = model.config.eos_token_id
     
     return f"Modèle {model_name} chargé avec succès sur GPU."
 
-def generate_text(input_text, temperature, top_p, top_k_value):
+def generate_text(input_text, temperature, top_p, top_k):
     global model, tokenizer
     
     inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
@@ -44,31 +45,25 @@ def generate_text(input_text, temperature, top_p, top_k_value):
             max_new_tokens=50,
             temperature=temperature,
             top_p=top_p,
-            top_k=top_k_value,
+            top_k=top_k,
             output_attentions=True,
-            output_scores=True,  # Activer les scores pour obtenir les logits
             return_dict_in_generate=True
         )
     
     generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
     
-    # Obtenir les logits pour le dernier token généré
+    # Logits et probabilités du dernier token généré
     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 5 des mots les plus probables
+    top_probs, top_indices = torch.topk(probabilities, 5)
     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
+    # Extraction des attentions
     attentions = torch.cat([att[-1].mean(dim=1) for att in outputs.attentions], dim=0).cpu().numpy()
-    
     attention_data = {
         'attention': attentions,
         'tokens': tokenizer.convert_ids_to_tokens(outputs.sequences[0])
@@ -107,6 +102,7 @@ def plot_probabilities(prob_data):
 def reset():
     return "", 1.0, 1.0, 50, None, None, None
 
+# Interface Gradio
 with gr.Blocks() as demo:
     gr.Markdown("# Générateur de texte avec visualisation d'attention")
     
@@ -131,6 +127,7 @@ with gr.Blocks() as demo:
     
     reset_button = gr.Button("Réinitialiser")
     
+    # Association des actions avec les boutons
     load_button.click(load_model, inputs=[model_dropdown], outputs=[load_output])
     generate_button.click(generate_text, 
                           inputs=[input_text, temperature, top_p, top_k], 
@@ -138,4 +135,5 @@ with gr.Blocks() as demo:
     reset_button.click(reset, 
                        outputs=[input_text, temperature, top_p, top_k, output_text, attention_plot, prob_plot])
 
+# Lancement de l'application
 demo.launch()