Update app.py

app.py

@@ -57,28 +57,13 @@ def load_model(model_name, progress=gr.Progress()):
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         progress(0.5, desc="Chargement du modèle")
         
-        # Configurations spécifiques par modèle
-        if "mixtral" in model_name.lower():
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                attn_implementation="flash_attention_2",
-                load_in_8bit=True
-            )
-        elif "llama" in model_name.lower() or "mistral" in model_name.lower():
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                device_map="auto",
-                attn_implementation="flash_attention_2"
-            )
-        else:
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                torch_dtype=torch.float16,
-                device_map="auto"
-            )
+        # Configuration générique pour tous les modèles
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            low_cpu_mem_usage=True
+        )
         
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token