Update app.py

app.py

@@ -9,11 +9,10 @@ import os
 # Login to Hugging Face with token
 login(token=os.environ["HF_TOKEN"])
 
-# Liste des modèles
-model_list = [
-    "meta-llama/Llama-2-13b",
-    "meta-llama/Llama-2-7b",
-    "meta-llama/Llama-2-70b",
+MODEL_LIST = [
+    "meta-llama/Llama-2-13b-hf",
+    "meta-llama/Llama-2-7b-hf",
+    "meta-llama/Llama-2-70b-hf",
     "meta-llama/Meta-Llama-3-8B",
     "meta-llama/Llama-3.2-3B",
     "meta-llama/Llama-3.1-8B",
@@ -26,126 +25,66 @@ model_list = [
     "croissantllm/CroissantLLMBase"
 ]
 
-# Charger le modèle et le tokenizer
-model = None
-tokenizer = None
+# Dictionnaire pour stocker les modèles et tokenizers déjà chargés
+loaded_models = {}
 
+# Charger le modèle
 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, attn_implementation="eager")
-    
-    print("Modèle chargé avec succès.")
-    return f"Modèle {model_name} chargé."
+    if model_name not in loaded_models:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+        model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+        loaded_models[model_name] = (model, tokenizer)
+    return loaded_models[model_name]
 
-def plot_attention(attention_data):
-    tokens = attention_data['tokens']
-    attention = attention_data['attention']
-    
-    fig, ax = plt.subplots(figsize=(10, 10))
-    cax = ax.matshow(attention, cmap='viridis')
-    fig.colorbar(cax)
+# Génération de texte et attention
+def generate_text(model_name, input_text, temperature, top_p, top_k):
+    model, tokenizer = load_model(model_name)
+    inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
 
-    ax.set_xticklabels([''] + tokens, rotation=90)
-    ax.set_yticklabels([''] + tokens)
+    # Génération du texte
+    output = model.generate(**inputs, max_new_tokens=50, temperature=temperature, top_p=top_p, top_k=top_k, output_attentions=True)
     
-    plt.xlabel("Tokens")
-    plt.ylabel("Tokens")
-    plt.title("Attention Heatmap")
-    
-    plt.tight_layout()
-    plt.savefig('attention_plot.png')
-    return 'attention_plot.png'
+    # Décodage de la sortie
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
-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
-    
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512).to(model.device)
-    
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=50,
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            output_scores=True,
-            output_attentions=True,
-            return_dict_in_generate=True,
-            return_legacy_cache=True
-        )
-    
-    generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
-    
-    # Logits et probabilités du dernier token généré
-    last_token_logits = outputs.scores[-1][0]
+    # Affichage des mots les plus probables
+    last_token_logits = output.scores[-1][0]
     probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
-    
-    # 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]
-    
-    prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
-    
-    # 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])
-    }
-    
-    return generated_text, plot_attention(attention_data), plot_probabilities(prob_data)
+    top_tokens = torch.topk(probabilities, k=5)
+    probable_words = [tokenizer.decode([token]) for token in top_tokens.indices]
 
-def reset_app():
-    global model, tokenizer
-    model = None
-    tokenizer = None
-    return "Application réinitialisée."
+    return generated_text, probable_words
 
 # Interface utilisateur Gradio
-with gr.Blocks() as demo:
-    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.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():
-        input_text = gr.Textbox(label="Entrer le texte")
-        generate_button = gr.Button("Générer")
-    
-    with gr.Row():
-        output_text = gr.Textbox(label="Texte généré", interactive=False)
-    
-    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():
-        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)
+def reset_interface():
+    return "", "", "", ""
+
+def main():
+    with gr.Blocks() as app:
+        with gr.Accordion("Choix du modèle", open=True):
+            model_name = gr.Dropdown(choices=MODEL_LIST, label="Modèles disponibles", value=MODEL_LIST[0])
+        
+        with gr.Row():
+            input_text = gr.Textbox(label="Texte d'entrée", placeholder="Saisissez votre texte ici...")
+        
+        with gr.Accordion("Paramètres", open=True):
+            temperature = gr.Slider(minimum=0, maximum=1, value=0.7, step=0.01, label="Température")
+            top_p = gr.Slider(minimum=0, maximum=1, value=0.9, step=0.01, label="Top_p")
+            top_k = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top_k")
+        
+        with gr.Row():
+            generate_button = gr.Button("Lancer la génération")
+            reset_button = gr.Button("Réinitialiser")
+        
+        generated_text_output = gr.Textbox(label="Texte généré", placeholder="Le texte généré s'affichera ici...")
+        probable_words_output = gr.Textbox(label="Mots les plus probables", placeholder="Les mots les plus probables apparaîtront ici...")
+        
+        # Lancer la génération
+        generate_button.click(generate_text, inputs=[model_name, input_text, temperature, top_p, top_k], outputs=[generated_text_output, probable_words_output])
+        # Réinitialiser
+        reset_button.click(reset_interface, outputs=[input_text, generated_text_output, probable_words_output])
+    
+    app.launch()
 
-demo.launch()
+if __name__ == "__main__":
+    main()