src/recommendations/loader.py

"""
Ce module contient des fonctions pour charger les modèles de recommandation et les données associées.
Il inclut des fonctions pour charger un modèle depuis le MLflow Model Registry, charger les données des cultures disponibles par zone à partir d'un fichier JSON, et charger un modèle à partir d'un dépôt Hugging Face.
"""
import json
import mlflow
import mlflow.pyfunc
from pathlib import Path
from huggingface_hub import snapshot_download
import joblib



def load_model_from_registry(
    model_name: str,
    tracking_uri: str,
    version: int = 1,
    stage: str | None = None
):
    """
    Charge un modèle depuis le MLflow Model Registry.
    Args:
        model_name (str): Le nom du modèle enregistré.
        tracking_uri (str): L'URI de suivi MLflow.
        version (int, optional): La version du modèle à charger. Par défaut à 1.
        stage (str, optional): Le stage du modèle à charger (ex: "Production"). Si spécifié, la version est ignorée.
    Returns:
        Le modèle chargé en tant qu'objet Python.
    """

    mlflow.set_tracking_uri(tracking_uri)
    mlflow.set_registry_uri(tracking_uri)

    if stage is not None:
        model_uri = f"models:/{model_name}/{stage}"
    else:
        model_uri = f"models:/{model_name}/{version}"

    model = mlflow.pyfunc.load_model(model_uri)
    return model


def load_crops_by_area(path):
    """Charge les données des cultures disponibles par zone à partir d'un fichier JSON.
    Args:
        path (str): Le chemin vers le fichier JSON contenant les données des cultures par zone.
    Returns:
        dict: Un dictionnaire contenant les cultures disponibles par zone.
    """
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


from pathlib import Path
import mlflow.pyfunc
from huggingface_hub import snapshot_download


def load_local_model(model_path):
    """
    Charge un modèle local depuis un dossier MLflow ou un fichier model.pkl.
    Args:
        model_path (str): Le chemin vers le dossier MLflow ou le fichier model.pkl.
    """
    model_path = Path(model_path)

    if model_path.is_dir() and (model_path / "model.pkl").exists():
        return joblib.load(model_path / "model.pkl")

    if model_path.is_file():
        return joblib.load(model_path)

    raise FileNotFoundError(f"Modèle introuvable : {model_path}")


def load_hf_model(repo_id, token=None):
    """
    Télécharge le modèle depuis Hugging Face et charge directement model.pkl.
    Args:
        repo_id (str): L'identifiant du dépôt Hugging Face contenant le modèle.
        token (str, optional): Le token d'authentification Hugging Face. Par défaut à None.
    Returns:
        Le modèle chargé en tant qu'objet Python.
    """
    local_dir = snapshot_download(
        repo_id=repo_id,
        repo_type="model",
        token=token
    )

    model_path = Path(local_dir) / "model.pkl"

    if not model_path.exists():
        raise FileNotFoundError(f"model.pkl introuvable dans {local_dir}")

    return joblib.load(model_path)


def load_model(model_source, local_path, hf_repo, hf_token=None):
    """Charge un modèle en fonction de la source spécifiée.
    Args:
        model_source (str): La source du modèle ("local" ou "hf").
        local_path (str): Le chemin vers le modèle local (utilisé si model_source est "local").
        hf_repo (str): L'identifiant du dépôt Hugging Face (utilisé si model_source est "hf").
        hf_token (str, optional): Le token d'authentification Hugging Face (utilisé si model_source est "hf"). Par défaut à None.
    Returns:
        Le modèle chargé en tant qu'objet Python.
    Raises:
        ValueError: Si model_source n'est pas "local" ou "hf".
    """
    model_source = model_source.strip().lower()

    if model_source == "local":
        return load_local_model(local_path)

    if model_source == "hf":
        return load_hf_model(hf_repo, hf_token)

    raise ValueError(f"MODEL_SOURCE invalide : {model_source}")