Update app.py

app.py

@@ -1,38 +1,75 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login
 import matplotlib.pyplot as plt
 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"])
+# Login to Hugging Face with token
+HF_TOKEN = "hf_token"  # Remplacer par ton token Hugging Face
+login(HF_TOKEN)
 
-# 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",
-    "mistralai/Mistral-7B-v0.1", "mistralai/Mixtral-8x7B-v0.1", "mistralai/Mistral-7B-v0.3",
-    "google/gemma-2-2b", "google/gemma-2-9b", "google/gemma-2-27b",
+# Liste des modèles
+model_list = [
+    "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",
+    "mistralai/Mistral-7B-v0.1",
+    "mistralai/Mixtral-8x7B-v0.1",
+    "mistralai/Mistral-7B-v0.3",
+    "google/gemma-2-2b",
+    "google/gemma-2-9b",
+    "google/gemma-2-27b",
     "croissantllm/CroissantLLMBase"
 ]
 
-# Variables pour le modèle et le tokenizer
+# Charger le modèle et le tokenizer
 model = None
 tokenizer = None
 
 def load_model(model_name):
     global model, tokenizer
+    print(f"Chargement du modèle {model_name}...")
+    
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16)
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype=torch.float16, attn_implementation="eager")
     
-    # 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
+    print("Modèle chargé avec succès.")
+    return f"Modèle {model_name} chargé."
+
+def plot_attention(attention_data):
+    tokens = attention_data['tokens']
+    attention = attention_data['attention']
     
-    return f"Modèle {model_name} chargé avec succès sur GPU."
+    fig, ax = plt.subplots(figsize=(10, 10))
+    cax = ax.matshow(attention, cmap='viridis')
+    fig.colorbar(cax)
+
+    ax.set_xticklabels([''] + tokens, rotation=90)
+    ax.set_yticklabels([''] + tokens)
+    
+    plt.xlabel("Tokens")
+    plt.ylabel("Tokens")
+    plt.title("Attention Heatmap")
+    
+    plt.tight_layout()
+    plt.savefig('attention_plot.png')
+    return 'attention_plot.png'
+
+def plot_probabilities(prob_data):
+    words, probs = zip(*prob_data.items())
+    
+    plt.figure(figsize=(6, 4))
+    plt.barh(words, probs, color='skyblue')
+    plt.xlabel('Probabilities')
+    plt.title('Top Probable Words')
+    
+    plt.tight_layout()
+    plt.savefig('probabilities_plot.png')
+    return 'probabilities_plot.png'
 
 def generate_text(input_text, temperature, top_p, top_k):
     global model, tokenizer
@@ -46,8 +83,10 @@ def generate_text(input_text, temperature, top_p, top_k):
             temperature=temperature,
             top_p=top_p,
             top_k=top_k,
+            output_scores=True,
             output_attentions=True,
-            return_dict_in_generate=True
+            return_dict_in_generate=True,
+            return_legacy_cache=True
         )
     
     generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
@@ -71,69 +110,42 @@ def generate_text(input_text, temperature, top_p, top_k):
     
     return generated_text, plot_attention(attention_data), plot_probabilities(prob_data)
 
-def plot_attention(attention_data):
-    attention = attention_data['attention']
-    tokens = attention_data['tokens']
-    
-    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())
-    
-    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)
-    plt.tight_layout()
-    return fig
-
-def reset():
-    return "", 1.0, 1.0, 50, None, None, None
+def reset_app():
+    global model, tokenizer
+    model = None
+    tokenizer = None
+    return "Application réinitialisée."
 
-# Interface Gradio
+# Interface utilisateur Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("# Générateur de texte avec visualisation d'attention")
+    with gr.Row():
+        model_selection = gr.Accordion("Sélection du modèle", open=True)
+        with model_selection:
+            model_name = gr.Dropdown(choices=model_list, label="Choisir un modèle", value=model_list[0])
+            load_model_button = gr.Button("Charger le modèle")
+            load_status = gr.Textbox(label="Statut du modèle", interactive=False)
     
-    with gr.Accordion("Sélection du modèle"):
-        model_dropdown = gr.Dropdown(choices=models, label="Choisissez un modèle")
-        load_button = gr.Button("Charger le modèle")
-        load_output = gr.Textbox(label="Statut du chargement")
+    with gr.Row():
+        temperature = gr.Slider(0.0, 1.0, value=0.7, label="Température")
+        top_p = gr.Slider(0.0, 1.0, value=0.9, label="Top-p")
+        top_k = gr.Slider(1, 100, value=50, label="Top-k")
     
     with gr.Row():
-        temperature = gr.Slider(0.1, 2.0, value=1.0, label="Température")
-        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="Entrer le texte")
+        generate_button = gr.Button("Générer")
     
-    input_text = gr.Textbox(label="Texte d'entrée")
-    generate_button = gr.Button("Générer")
+    with gr.Row():
+        output_text = gr.Textbox(label="Texte généré", interactive=False)
     
-    output_text = gr.Textbox(label="Texte généré")
+    with gr.Row():
+        attention_plot = gr.Image(label="Carte de chaleur des attentions")
+        prob_plot = gr.Image(label="Probabilités des mots les plus probables")
     
     with gr.Row():
-        attention_plot = gr.Plot(label="Visualisation de l'attention")
-        prob_plot = gr.Plot(label="Probabilités des tokens suivants")
-    
-    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], 
-                          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])
+        reset_button = gr.Button("Réinitialiser l'application")
+    
+    load_model_button.click(load_model, inputs=[model_name], outputs=[load_status])
+    generate_button.click(generate_text, inputs=[input_text, temperature, top_p, top_k], outputs=[output_text, attention_plot, prob_plot])
+    reset_button.click(reset_app)
 
-# Lancement de l'application
 demo.launch()