llm_ticket3/tests/common/llm_factory.py
2025-04-18 17:34:21 +02:00

114 lines
3.5 KiB
Python

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Factory pour les modèles LLM utilisés dans les tests.
Centralise l'initialisation des modèles pour éviter la duplication de code.
"""
import logging
import time
from typing import Dict, Any, Optional, Type, List, Tuple
# Import des modèles LLM disponibles
from llm_classes.mistral_large import MistralLarge
from llm_classes.mistral_medium import MistralMedium
from llm_classes.pixtral_large import PixtralLarge
from llm_classes.pixtral_12b import Pixtral12b
from llm_classes.deepseek import DeepSeek
from llm_classes.llama_vision import LlamaVision
from llm_classes.qwen2_5 import Qwen2_5
logger = logging.getLogger("LLMFactory")
# Mapping des noms de modèles aux classes
LLM_CLASSES = {
"mistral_large": MistralLarge,
"mistral_medium": MistralMedium,
"pixtral_large": PixtralLarge,
"pixtral12b": Pixtral12b,
"deepseek": DeepSeek,
"llama_vision": LlamaVision,
"qwen2_5": Qwen2_5
}
# Catégories de modèles par capacité
TEXT_MODELS = ["mistral_large", "mistral_medium", "deepseek", "qwen2_5"]
VISION_MODELS = ["pixtral_large", "pixtral12b", "llama_vision"]
def create_llm(model_name: str, **kwargs) -> Any:
"""
Crée une instance d'un modèle LLM spécifique.
Args:
model_name: Nom du modèle à créer
**kwargs: Arguments supplémentaires à passer au constructeur
Returns:
Instance du modèle LLM demandé
Raises:
ValueError: Si le modèle demandé n'existe pas
"""
if model_name not in LLM_CLASSES:
raise ValueError(f"Modèle LLM inconnu: {model_name}. Options disponibles: {', '.join(LLM_CLASSES.keys())}")
logger.info(f"Initialisation du modèle {model_name}")
start_time = time.time()
model = LLM_CLASSES[model_name](**kwargs)
init_time = time.time() - start_time
logger.info(f"Modèle {model_name} initialisé en {init_time:.2f} secondes")
return model
def create_multiple_llms(model_names: List[str], **kwargs) -> Dict[str, Any]:
"""
Crée plusieurs instances de modèles LLM.
Args:
model_names: Liste des noms de modèles à créer
**kwargs: Arguments supplémentaires à passer aux constructeurs
Returns:
Dictionnaire avec les noms de modèles en clés et les instances en valeurs
"""
models = {}
for model_name in model_names:
models[model_name] = create_llm(model_name, **kwargs)
return models
def get_best_text_model() -> Any:
"""
Retourne le meilleur modèle texte disponible.
Returns:
Instance du meilleur modèle texte
"""
for model_name in ["mistral_large", "qwen2_5", "deepseek", "mistral_medium"]:
try:
return create_llm(model_name)
except Exception as e:
logger.warning(f"Impossible d'initialiser {model_name}: {e}")
# Si aucun modèle n'a pu être initialisé
raise ValueError("Aucun modèle texte n'a pu être initialisé")
def get_best_vision_model() -> Any:
"""
Retourne le meilleur modèle multimodal (vision) disponible.
Returns:
Instance du meilleur modèle vision
"""
for model_name in ["pixtral_large", "pixtral12b", "llama_vision"]:
try:
return create_llm(model_name)
except Exception as e:
logger.warning(f"Impossible d'initialiser {model_name}: {e}")
# Si aucun modèle n'a pu être initialisé
raise ValueError("Aucun modèle vision n'a pu être initialisé")