app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import login
import os
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
import time

# Authentification
login(token=os.environ["HF_TOKEN"])

# Liste des modèles
models = [
    "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",
    "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 globales
model = None
tokenizer = None

def load_model(model_name, progress=gr.Progress()):
    global model, tokenizer
    try:
        for i in progress.tqdm(range(100)):
            time.sleep(0.01)  # Simuler le chargement
            if i == 25:
                tokenizer = AutoTokenizer.from_pretrained(model_name)
            elif i == 75:
                model = AutoModelForCausalLM.from_pretrained(
                    model_name, 
                    torch_dtype=torch.float32,
                    device_map="cpu",
                    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."
    except Exception as e:
        return f"Erreur lors du chargement du modèle : {str(e)}"

def analyze_next_token(input_text, temperature, top_p, top_k):
    global model, tokenizer
    
    if model is None or tokenizer is None:
        return "Veuillez d'abord charger un modèle.", None, None

    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    
    try:
        with torch.no_grad():
            outputs = model(**inputs, output_attentions=True)
        
        last_token_logits = outputs.logits[0, -1, :]
        probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
        
        # Obtenir les 10 tokens les plus probables
        top_k = 10
        top_probs, top_indices = torch.topk(probabilities, top_k)
        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)}
        
        # Créer un texte explicatif
        prob_text = "Prochains tokens les plus probables :\n\n"
        for word, prob in prob_data.items():
            escaped_word = word.replace("<", "&lt;").replace(">", "&gt;")
            prob_text += f"{escaped_word}: {prob:.2%}\n"
        
        # Créer les visualisations
        prob_plot = plot_probabilities(prob_data)
        attention_plot = plot_attention(inputs["input_ids"][0], outputs.attentions)
        
        return prob_text, attention_plot, prob_plot
    except Exception as e:
        return f"Erreur lors de l'analyse : {str(e)}", None, None

def generate_text(input_text, temperature, top_p, top_k):
    global model, tokenizer
    
    if model is None or tokenizer is None:
        return "Veuillez d'abord charger un modèle."

    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)
    
    try:
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=1,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k
            )
        
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return generated_text  # Retourne l'input + le nouveau mot
    except Exception as e:
        return f"Erreur lors de la génération : {str(e)}"

def plot_probabilities(prob_data):
    words = list(prob_data.keys())
    probs = list(prob_data.values())
    
    fig, ax = plt.subplots(figsize=(12, 6))
    bars = ax.bar(range(len(words)), probs, color='lightgreen')
    ax.set_title("Probabilités des tokens suivants les plus probables")
    ax.set_xlabel("Tokens")
    ax.set_ylabel("Probabilité")
    
    ax.set_xticks(range(len(words)))
    ax.set_xticklabels(words, rotation=45, ha='right')
    
    for i, (bar, word) in enumerate(zip(bars, words)):
        height = bar.get_height()
        ax.text(i, height, f'{word}\n{height:.2%}',
                ha='center', va='bottom', rotation=0)
    
    plt.tight_layout()
    return fig

def plot_attention(input_ids, attention_outputs):
    input_tokens = tokenizer.convert_ids_to_tokens(input_ids)
    
    # Prendre la moyenne des attentions sur toutes les couches et têtes
    attention = torch.mean(torch.cat(attention_outputs), dim=(0, 1)).cpu().numpy()
    
    fig, ax = plt.subplots(figsize=(12, 10))
    sns.heatmap(attention, annot=True, cmap="YlOrRd", xticklabels=input_tokens, yticklabels=input_tokens, ax=ax)
    
    ax.set_title("Carte d'attention moyenne")
    plt.tight_layout()
    return fig

def reset():
    global model, tokenizer
    model = None
    tokenizer = None
    return "", 1.0, 1.0, 50, None, None, None, None

with gr.Blocks() as demo:
    gr.Markdown("# Analyse et génération de texte")
    
    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.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="Texte d'entrée", lines=3)
    analyze_button = gr.Button("Analyser le prochain token")
    
    next_token_probs = gr.Textbox(label="Probabilités du prochain token")
    
    with gr.Row():
        attention_plot = gr.Plot(label="Visualisation de l'attention")
        prob_plot = gr.Plot(label="Probabilités des tokens suivants")
    
    generate_button = gr.Button("Générer le prochain mot")
    generated_text = gr.Textbox(label="Texte 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])
    generate_button.click(generate_text, 
                          inputs=[input_text, temperature, top_p, top_k], 
                          outputs=[generated_text])
    reset_button.click(reset, 
                       outputs=[input_text, temperature, top_p, top_k, next_token_probs, attention_plot, prob_plot, generated_text])

if __name__ == "__main__":
    demo.launch()