Update app.py

app.py

@@ -12,54 +12,54 @@ import time
 login(token=os.environ["HF_TOKEN"])
 
 # Structure hiérarchique des modèles
-model_hierarchy = {
+models_hierarchy = {
     "meta-llama": {
-        "Llama-2": ["7B", "13B", "70B"],
-        "Llama-3": ["8B", "3.2B", "3.1B"]
+        "Llama-2": ["7b", "13b", "70b"],
+        "Llama-3": ["8B", "3.2-3B", "3.1-8B"]
     },
     "mistralai": {
         "Mistral": ["7B-v0.1", "7B-v0.3"],
         "Mixtral": ["8x7B-v0.1"]
     },
     "google": {
-        "Gemma": ["2B", "9B", "27B"]
+        "gemma": ["2b", "9b", "27b"]
     },
     "croissantllm": {
         "CroissantLLM": ["Base"]
     }
 }
 
-# Mise à jour de la liste des modèles et leurs langues supportées
+# Langues supportées par modèle
 models_and_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"],
+    "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-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}
 }
 
@@ -69,14 +69,17 @@ tokenizer = None
 selected_language = None
 
 def update_model_choices(company):
-    return gr.Dropdown(choices=list(model_hierarchy[company].keys()), value=None)
+    return gr.Dropdown(choices=list(models_hierarchy[company].keys()), value=None)
 
 def update_variation_choices(company, model_name):
-    return gr.Dropdown(choices=model_hierarchy[company][model_name], value=None)
+    return gr.Dropdown(choices=models_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")
@@ -116,14 +119,13 @@ 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é
-# ...
+# Le reste du code reste inchangé...
 
 with gr.Blocks() as demo:
     gr.Markdown("# LLM&BIAS")
     
     with gr.Accordion("Sélection du modèle"):
-        company_dropdown = gr.Dropdown(choices=list(model_hierarchy.keys()), label="Choisissez une société")
+        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=[])
         load_button = gr.Button("Charger le modèle")