Update app.py

app.py

@@ -12,9 +12,9 @@ import time
 login(token=os.environ["HF_TOKEN"])
 
 # Structure hiérarchique des modèles
-models_hierarchy = {
+model_hierarchy = {
     "meta-llama": {
-        "Llama-2": ["7b", "13b", "70b"],
+        "Llama-2": ["7B", "13B", "70B"],
         "Llama-3": ["8B", "3.2-3B", "3.1-8B"]
     },
     "mistralai": {
@@ -22,7 +22,7 @@ models_hierarchy = {
         "Mixtral": ["8x7B-v0.1"]
     },
     "google": {
-        "gemma": ["2b", "9b", "27b"]
+        "Gemma": ["2B", "9B", "27B"]
     },
     "croissantllm": {
         "CroissantLLM": ["Base"]
@@ -31,35 +31,35 @@ models_hierarchy = {
 
 # Langues supportées par modèle
 models_and_languages = {
-    "meta-llama/Llama-2-7b-hf": ["en"],
-    "meta-llama/Llama-2-13b-hf": ["en"],
-    "meta-llama/Llama-2-70b-hf": ["en"],
-    "meta-llama/Meta-Llama-3-8B": ["en"],
+    "meta-llama/Llama-2-7B": ["en"],
+    "meta-llama/Llama-2-13B": ["en"],
+    "meta-llama/Llama-2-70B": ["en"],
+    "meta-llama/Llama-3-8B": ["en"],
     "meta-llama/Llama-3.2-3B": ["en", "de", "fr", "it", "pt", "hi", "es", "th"],
     "meta-llama/Llama-3.1-8B": ["en", "de", "fr", "it", "pt", "hi", "es", "th"],
     "mistralai/Mistral-7B-v0.1": ["en"],
-    "mistralai/Mixtral-8x7B-v0.1": ["en", "fr", "it", "de", "es"],
     "mistralai/Mistral-7B-v0.3": ["en"],
-    "google/gemma-2-2b": ["en"],
-    "google/gemma-2-9b": ["en"],
-    "google/gemma-2-27b": ["en"],
+    "mistralai/Mixtral-8x7B-v0.1": ["en", "fr", "it", "de", "es"],
+    "google/Gemma-2B": ["en"],
+    "google/Gemma-9B": ["en"],
+    "google/Gemma-27B": ["en"],
     "croissantllm/CroissantLLMBase": ["en", "fr"]
 }
 
 # Paramètres recommandés pour chaque modèle
 model_parameters = {
-    "meta-llama/Llama-2-13b-hf": {"temperature": 0.8, "top_p": 0.9, "top_k": 40},
-    "meta-llama/Llama-2-7b-hf": {"temperature": 0.8, "top_p": 0.9, "top_k": 40},
-    "meta-llama/Llama-2-70b-hf": {"temperature": 0.8, "top_p": 0.9, "top_k": 40},
-    "meta-llama/Meta-Llama-3-8B": {"temperature": 0.75, "top_p": 0.9, "top_k": 50},
+    "meta-llama/Llama-2-7B": {"temperature": 0.8, "top_p": 0.9, "top_k": 40},
+    "meta-llama/Llama-2-13B": {"temperature": 0.8, "top_p": 0.9, "top_k": 40},
+    "meta-llama/Llama-2-70B": {"temperature": 0.8, "top_p": 0.9, "top_k": 40},
+    "meta-llama/Llama-3-8B": {"temperature": 0.75, "top_p": 0.9, "top_k": 50},
     "meta-llama/Llama-3.2-3B": {"temperature": 0.75, "top_p": 0.9, "top_k": 50},
     "meta-llama/Llama-3.1-8B": {"temperature": 0.75, "top_p": 0.9, "top_k": 50},
     "mistralai/Mistral-7B-v0.1": {"temperature": 0.7, "top_p": 0.9, "top_k": 50},
-    "mistralai/Mixtral-8x7B-v0.1": {"temperature": 0.8, "top_p": 0.95, "top_k": 50},
     "mistralai/Mistral-7B-v0.3": {"temperature": 0.7, "top_p": 0.9, "top_k": 50},
-    "google/gemma-2-2b": {"temperature": 0.7, "top_p": 0.95, "top_k": 40},
-    "google/gemma-2-9b": {"temperature": 0.7, "top_p": 0.95, "top_k": 40},
-    "google/gemma-2-27b": {"temperature": 0.7, "top_p": 0.95, "top_k": 40},
+    "mistralai/Mixtral-8x7B-v0.1": {"temperature": 0.8, "top_p": 0.95, "top_k": 50},
+    "google/Gemma-2B": {"temperature": 0.7, "top_p": 0.95, "top_k": 40},
+    "google/Gemma-9B": {"temperature": 0.7, "top_p": 0.95, "top_k": 40},
+    "google/Gemma-27B": {"temperature": 0.7, "top_p": 0.95, "top_k": 40},
     "croissantllm/CroissantLLMBase": {"temperature": 0.8, "top_p": 0.92, "top_k": 50}
 }
 
@@ -69,24 +69,20 @@ tokenizer = None
 selected_language = None
 
 def update_model_choices(company):
-    return gr.Dropdown(choices=list(models_hierarchy[company].keys()), value=None)
+    return gr.Dropdown(choices=list(model_hierarchy[company].keys()), value=None)
 
 def update_variation_choices(company, model_name):
-    return gr.Dropdown(choices=models_hierarchy[company][model_name], value=None)
+    return gr.Dropdown(choices=model_hierarchy[company][model_name], value=None)
 
 def load_model(company, model_name, variation, progress=gr.Progress()):
     global model, tokenizer
-    
     full_model_name = f"{company}/{model_name}-{variation}"
-    if full_model_name not in models_and_languages:
-        full_model_name = f"{company}/{model_name}{variation}"
     
     try:
         progress(0, desc="Chargement du tokenizer")
         tokenizer = AutoTokenizer.from_pretrained(full_model_name)
         progress(0.5, desc="Chargement du modèle")
         
-        # Configurations spécifiques par modèle
         if "mixtral" in full_model_name.lower():
             model = AutoModelForCausalLM.from_pretrained(
                 full_model_name,
@@ -106,9 +102,8 @@ def load_model(company, model_name, variation, progress=gr.Progress()):
         
         progress(1.0, desc="Modèle chargé")
         available_languages = models_and_languages[full_model_name]
-        
-        # Mise à jour des sliders avec les valeurs recommandées
         params = model_parameters[full_model_name]
+        
         return (
             f"Modèle {full_model_name} chargé avec succès. Langues disponibles : {', '.join(available_languages)}",
             gr.Dropdown(choices=available_languages, value=available_languages[0], visible=True, interactive=True),
@@ -119,15 +114,129 @@ def load_model(company, model_name, variation, progress=gr.Progress()):
     except Exception as e:
         return f"Erreur lors du chargement du modèle : {str(e)}", gr.Dropdown(visible=False), None, None, None
 
-# Le reste du code reste inchangé...
+def set_language(lang):
+    global selected_language
+    selected_language = lang
+    return f"Langue sélectionnée : {lang}"
+
+def ensure_token_display(token):
+    if token.isdigit() or (token.startswith('-') and token[1:].isdigit()):
+        return tokenizer.decode([int(token)])
+    return token
+
+def analyze_next_token(input_text, temperature, top_p, top_k):
+    global model, tokenizer, selected_language
+    
+    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).to(model.device)
+    
+    try:
+        with torch.no_grad():
+            outputs = model(**inputs)
+        
+        last_token_logits = outputs.logits[0, -1, :]
+        probabilities = torch.nn.functional.softmax(last_token_logits, dim=-1)
+        
+        top_k = 10
+        top_probs, top_indices = torch.topk(probabilities, top_k)
+        top_words = [ensure_token_display(tokenizer.decode([idx.item()])) for idx in top_indices]
+        prob_data = {word: prob.item() for word, prob in zip(top_words, top_probs)}
+        
+        prob_text = "Prochains tokens les plus probables :\n\n"
+        for word, prob in prob_data.items():
+            prob_text += f"{word}: {prob:.2%}\n"
+        
+        prob_plot = plot_probabilities(prob_data)
+        attention_plot = plot_attention(inputs["input_ids"][0].cpu(), last_token_logits.cpu())
+        
+        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, selected_language
+    
+    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).to(model.device)
+    
+    try:
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=10,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k
+            )
+        
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return generated_text
+    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'{height:.2%}',
+                ha='center', va='bottom', rotation=0)
+    
+    plt.tight_layout()
+    return fig
+
+def plot_attention(input_ids, last_token_logits):
+    input_tokens = [ensure_token_display(tokenizer.decode([id])) for id in input_ids]
+    attention_scores = torch.nn.functional.softmax(last_token_logits, dim=-1)
+    top_k = min(len(input_tokens), 10)
+    top_attention_scores, _ = torch.topk(attention_scores, top_k)
+    
+    fig, ax = plt.subplots(figsize=(14, 7))
+    sns.heatmap(top_attention_scores.unsqueeze(0).numpy(), annot=True, cmap="YlOrRd", cbar=True, ax=ax, fmt='.2%')
+    ax.set_xticklabels(input_tokens[-top_k:], rotation=45, ha="right", fontsize=10)
+    ax.set_yticklabels(["Attention"], rotation=0, fontsize=10)
+    ax.set_title("Scores d'attention pour les derniers tokens", fontsize=16)
+    
+    cbar = ax.collections[0].colorbar
+    cbar.set_label("Score d'attention", fontsize=12)
+    cbar.ax.tick_params(labelsize=10)
+    
+    plt.tight_layout()
+    return fig
+
+def reset():
+    global model, tokenizer, selected_language
+    model = None
+    tokenizer = None
+    selected_language = None
+    return (
+        gr.Dropdown(choices=list(model_hierarchy.keys()), value=None),
+        gr.Dropdown(visible=False),
+        gr.Dropdown(visible=False),
+        "", 1.0, 1.0, 50, None, None, None, None, gr.Dropdown(visible=False), ""
+    )
 
 with gr.Blocks() as demo:
     gr.Markdown("# LLM&BIAS")
     
     with gr.Accordion("Sélection du modèle"):
-        company_dropdown = gr.Dropdown(choices=list(models_hierarchy.keys()), label="Choisissez une société")
-        model_dropdown = gr.Dropdown(label="Choisissez un modèle", choices=[])
-        variation_dropdown = gr.Dropdown(label="Choisissez une variation", choices=[])
+        company_dropdown = gr.Dropdown(choices=list(model_hierarchy.keys()), label="Choisissez une société")
+        model_dropdown = gr.Dropdown(label="Choisissez un modèle", visible=False)
+        variation_dropdown = gr.Dropdown(label="Choisissez une variation", visible=False)
         load_button = gr.Button("Charger le modèle")
         load_output = gr.Textbox(label="Statut du chargement")
         language_dropdown = gr.Dropdown(label="Choisissez une langue", visible=False)
@@ -156,7 +265,7 @@ with gr.Blocks() as demo:
     model_dropdown.change(update_variation_choices, inputs=[company_dropdown, model_dropdown], outputs=[variation_dropdown])
     load_button.click(load_model, 
                       inputs=[company_dropdown, model_dropdown, variation_dropdown], 
-                      outputs=[load_output, language_dropdown, temperature, top_p, top_k])
+                      outputs=[load_output, language_dropdown])
     language_dropdown.change(set_language, inputs=[language_dropdown], outputs=[language_output])
     analyze_button.click(analyze_next_token, 
                          inputs=[input_text, temperature, top_p, top_k],