Update app.py

app.py

@@ -11,21 +11,41 @@ import time
 # Authentification
 login(token=os.environ["HF_TOKEN"])
 
-# Liste des modèles et leurs langues supportées
-models_and_languages = {
-    "meta-llama/Llama-2-13b-hf": ["en"],
-    "meta-llama/Llama-2-7b-hf": ["en"],
-    "meta-llama/Llama-2-70b-hf": ["en"],
-    "meta-llama/Meta-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"],
-    "croissantllm/CroissantLLMBase": ["en", "fr"]
+# Restructuration des modèles et de leurs informations
+models_info = {
+    "Meta-llama": {
+        "Llama 2": {
+            "7B": {"name": "meta-llama/Llama-2-7b-hf", "languages": ["en"]},
+            "13B": {"name": "meta-llama/Llama-2-13b-hf", "languages": ["en"]},
+            "70B": {"name": "meta-llama/Llama-2-70b-hf", "languages": ["en"]},
+        },
+        "Llama 3": {
+            "8B": {"name": "meta-llama/Meta-Llama-3-8B", "languages": ["en"]},
+            "3.2-3B": {"name": "meta-llama/Llama-3.2-3B", "languages": ["en", "de", "fr", "it", "pt", "hi", "es", "th"]},
+            "3.1-8B": {"name": "meta-llama/Llama-3.1-8B", "languages": ["en", "de", "fr", "it", "pt", "hi", "es", "th"]},
+        },
+    },
+    "Mistral AI": {
+        "Mistral": {
+            "7B-v0.1": {"name": "mistralai/Mistral-7B-v0.1", "languages": ["en"]},
+            "7B-v0.3": {"name": "mistralai/Mistral-7B-v0.3", "languages": ["en"]},
+        },
+        "Mixtral": {
+            "8x7B-v0.1": {"name": "mistralai/Mixtral-8x7B-v0.1", "languages": ["en", "fr", "it", "de", "es"]},
+        },
+    },
+    "Google": {
+        "Gemma": {
+            "2B": {"name": "google/gemma-2-2b", "languages": ["en"]},
+            "9B": {"name": "google/gemma-2-9b", "languages": ["en"]},
+            "27B": {"name": "google/gemma-2-27b", "languages": ["en"]},
+        },
+    },
+    "CroissantLLM": {
+        "CroissantLLMBase": {
+            "Base": {"name": "croissantllm/CroissantLLMBase", "languages": ["en", "fr"]},
+        },
+    },
 }
 
 # Paramètres recommandés pour chaque modèle
@@ -50,6 +70,18 @@ model = None
 tokenizer = None
 selected_language = None
 
+def update_model_type(family):
+    return gr.Dropdown(choices=list(models_info[family].keys()), value=None, interactive=True)
+
+def update_model_variation(family, model_type):
+    return gr.Dropdown(choices=list(models_info[family][model_type].keys()), value=None, interactive=True)
+
+def update_selected_model(family, model_type, variation):
+    if family and model_type and variation:
+        model_name = models_info[family][model_type][variation]["name"]
+        return model_name, gr.Dropdown(choices=models_info[family][model_type][variation]["languages"], value=models_info[family][model_type][variation]["languages"][0], visible=True, interactive=True)
+    return "", gr.Dropdown(visible=False)
+
 def load_model(model_name, progress=gr.Progress()):
     global model, tokenizer
     try:
@@ -76,7 +108,15 @@ def load_model(model_name, progress=gr.Progress()):
             tokenizer.pad_token = tokenizer.eos_token
         
         progress(1.0, desc="Modèle chargé")
-        available_languages = models_and_languages[model_name]
+        
+        # Recherche des langues disponibles pour le modèle sélectionné
+        available_languages = next(
+            (info["languages"] for family in models_info.values()
+             for model_type in family.values()
+             for variation in model_type.values()
+             if variation["name"] == model_name),
+            ["en"]  # Défaut à l'anglais si non trouvé
+        )
         
         # Mise à jour des sliders avec les valeurs recommandées
         params = model_parameters[model_name]
@@ -200,13 +240,24 @@ def reset():
     model = None
     tokenizer = None
     selected_language = None
-    return "", 1.0, 1.0, 50, None, None, None, None, gr.Dropdown(visible=False), ""
+    return (
+        "", 1.0, 1.0, 50, None, None, None, None,
+        gr.Dropdown(choices=list(models_info.keys()), value=None, interactive=True),
+        gr.Dropdown(choices=[], value=None, interactive=False),
+        gr.Dropdown(choices=[], value=None, interactive=False),
+        "", gr.Dropdown(visible=False), ""
+    )
 
 with gr.Blocks() as demo:
     gr.Markdown("# LLM&BIAS")
     
-    with gr.Accordion("Sélection du modèle"):
-        model_dropdown = gr.Dropdown(choices=list(models_and_languages.keys()), label="Choisissez un modèle")
+    with gr.Accordion("Sélection du modèle", open=True):
+        with gr.Row():
+            model_family = gr.Dropdown(choices=list(models_info.keys()), label="Famille de modèle", interactive=True)
+            model_type = gr.Dropdown(choices=[], label="Type de modèle", interactive=False)
+            model_variation = gr.Dropdown(choices=[], label="Variation du modèle", interactive=False)
+        
+        selected_model = gr.Textbox(label="Modèle sélectionné", interactive=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)
@@ -231,18 +282,58 @@ with gr.Blocks() as demo:
     
     reset_button = gr.Button("Réinitialiser")
     
-    load_button.click(load_model, 
-                      inputs=[model_dropdown], 
-                      outputs=[load_output, language_dropdown, temperature, top_p, top_k])
-    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], 
-                         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, language_dropdown, language_output])
+    # Événements pour la sélection du modèle
+    model_family.change(
+        update_model_type,
+        inputs=[model_family],
+        outputs=[model_type]
+    )
+    
+    model_type.change(
+        update_model_variation,
+        inputs=[model_family, model_type],
+        outputs=[model_variation]
+    )
+    
+    model_variation.change(
+        update_selected_model,
+        inputs=[model_family, model_type, model_variation],
+        outputs=[selected_model, language_dropdown]
+    )
+    
+    load_button.click(
+        load_model,
+        inputs=[selected_model],
+        outputs=[load_output, language_dropdown, temperature, top_p, top_k]
+    )
+    
+    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],
+        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,
+            model_family, model_type, model_variation,
+            selected_model, language_dropdown, language_output
+        ]
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()