Update app.py

app.py

@@ -11,55 +11,37 @@ import time
 # Authentification
 login(token=os.environ["HF_TOKEN"])
 
-# Structure hiérarchique des modèles
-models_hierarchy = {
-    "meta-llama": {
-        "Llama-2": ["7B", "13B", "70B"],
-        "Llama-3": ["8B", "3.2B", "3.1B"]
-    },
-    "mistralai": {
-        "Mistral": ["7B-v0.1", "7B-v0.3"],
-        "Mixtral": ["8x7B-v0.1"]
-    },
-    "google": {
-        "Gemma": ["2B", "9B", "27B"]
-    },
-    "croissantllm": {
-        "CroissantLLM": ["Base"]
-    }
-}
-
-# Langues supportées par modèle
-models_languages = {
-    "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-3.2B": ["en", "de", "fr", "it", "pt", "hi", "es", "th"],
-    "meta-llama/Llama-3-3.1B": ["en", "de", "fr", "it", "pt", "hi", "es", "th"],
+# 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-2B": ["en"],
-    "google/Gemma-9B": ["en"],
-    "google/Gemma-27B": ["en"],
+    "google/gemma-2-2b": ["en"],
+    "google/gemma-2-9b": ["en"],
+    "google/gemma-2-27b": ["en"],
     "croissantllm/CroissantLLMBase": ["en", "fr"]
 }
 
 # Paramètres recommandés pour chaque modèle
 model_parameters = {
-    "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-3.2B": {"temperature": 0.75, "top_p": 0.9, "top_k": 50},
-    "meta-llama/Llama-3-3.1B": {"temperature": 0.75, "top_p": 0.9, "top_k": 50},
+    "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-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-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},
+    "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},
     "croissantllm/CroissantLLMBase": {"temperature": 0.8, "top_p": 0.92, "top_k": 50}
 }
 
@@ -68,31 +50,24 @@ model = None
 tokenizer = None
 selected_language = None
 
-def update_model_choices(company):
-    return list(models_hierarchy[company].keys())
-
-def update_variation_choices(company, model_name):
-    return models_hierarchy[company][model_name]
-
-def load_model(company, model_name, variation, progress=gr.Progress()):
+def load_model(model_name, progress=gr.Progress()):
     global model, tokenizer
-    full_model_name = f"{company}/{model_name}-{variation}"
-    
     try:
         progress(0, desc="Chargement du tokenizer")
-        tokenizer = AutoTokenizer.from_pretrained(full_model_name)
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
         progress(0.5, desc="Chargement du modèle")
         
-        if "mixtral" in full_model_name.lower():
+        # Configurations spécifiques par modèle
+        if "mixtral" in model_name.lower():
             model = AutoModelForCausalLM.from_pretrained(
-                full_model_name,
+                model_name,
                 torch_dtype=torch.float16,
                 device_map="auto",
                 load_in_8bit=True
             )
         else:
             model = AutoModelForCausalLM.from_pretrained(
-                full_model_name,
+                model_name,
                 torch_dtype=torch.float16,
                 device_map="auto"
             )
@@ -101,11 +76,12 @@ def load_model(company, model_name, variation, progress=gr.Progress()):
             tokenizer.pad_token = tokenizer.eos_token
         
         progress(1.0, desc="Modèle chargé")
-        available_languages = models_languages[full_model_name]
+        available_languages = models_and_languages[model_name]
         
-        params = model_parameters[full_model_name]
+        # Mise à jour des sliders avec les valeurs recommandées
+        params = model_parameters[model_name]
         return (
-            f"Modèle {full_model_name} chargé avec succès. Langues disponibles : {', '.join(available_languages)}",
+            f"Modèle {model_name} chargé avec succès. Langues disponibles : {', '.join(available_languages)}",
             gr.Dropdown(choices=available_languages, value=available_languages[0], visible=True, interactive=True),
             params["temperature"],
             params["top_p"],
@@ -120,6 +96,7 @@ def set_language(lang):
     return f"Langue sélectionnée : {lang}"
 
 def ensure_token_display(token):
+    """Assure que le token est affiché correctement."""
     if token.isdigit() or (token.startswith('-') and token[1:].isdigit()):
         return tokenizer.decode([int(token)])
     return token
@@ -229,9 +206,7 @@ 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", interactive=False)
-        variation_dropdown = gr.Dropdown(label="Choisissez une variation", interactive=False)
+        model_dropdown = gr.Dropdown(choices=list(models_and_languages.keys()), label="Choisissez un modèle")
         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)
@@ -256,11 +231,8 @@ with gr.Blocks() as demo:
     
     reset_button = gr.Button("Réinitialiser")
     
-    company_dropdown.change(update_model_choices, inputs=[company_dropdown], outputs=[model_dropdown])
-    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], 
+                      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, 
@@ -273,4 +245,4 @@ with gr.Blocks() as demo:
                        outputs=[input_text, temperature, top_p, top_k, next_token_probs, attention_plot, prob_plot, generated_text, language_dropdown, language_output])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()