Skip to content

# Module: `model_loader`

Model Loader: depending on the model source, load the model from the local file system, Hugging Face Hub, or Arweave.

## `ArweaveLoadArgs`

Bases: `CommonLoadArgs`

Arguments for loading the model

Source code in `src/infernet_ml/utils/model_loader.py`

    54
    55
    56
    57
    58
    59
    class ArweaveLoadArgs(CommonLoadArgs):
        """
        Arguments for loading the model
        """

        pass

## `CommonLoadArgs`

Bases: `BaseModel`

Common arguments for loading a model

Source code in `src/infernet_ml/utils/model_loader.py`

    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    class CommonLoadArgs(BaseModel):
        """
        Common arguments for loading a model
        """

        model_config = ConfigDict(frozen=True)

        cache_path: Optional[str] = None
        version: Optional[str] = None
        repo_id: str
        filename: str

## `HFLoadArgs`

Bases: `CommonLoadArgs`

Arguments for loading the model

Source code in `src/infernet_ml/utils/model_loader.py`

    46
    47
    48
    49
    50
    51
    class HFLoadArgs(CommonLoadArgs):
        """
        Arguments for loading the model
        """

        pass

## `LocalLoadArgs`

Bases: `BaseModel`

Arguments for loading the model

Source code in `src/infernet_ml/utils/model_loader.py`

    62
    63
    64
    65
    66
    67
    68
    69
    class LocalLoadArgs(BaseModel):
        """
        Arguments for loading the model
        """

        model_config = ConfigDict(frozen=True)

        path: str

## `ModelSource`

Bases: `IntEnum`

Enum for the model source

- `LOCAL`: Load the model from the local file system
- `ARWEAVE`: Load the model from Arweave
- `HUGGINGFACE_HUB`: Load the model from Hugging Face Hub

Source code in `src/infernet_ml/utils/model_loader.py`

    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    class ModelSource(IntEnum):
        """
        Enum for the model source

        - `LOCAL`: Load the model from the local file system
        - `ARWEAVE`: Load the model from Arweave
        - `HUGGINGFACE_HUB`: Load the model from Hugging Face Hub
        """

        LOCAL = 0
        ARWEAVE = 1
        HUGGINGFACE_HUB = 2

## `download_model(model_source, load_args)`

Load the model from the specified source.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`model_source` | `ModelSource` | the source of the model | required
`load_args` | `LoadArgs` | the load arguments, options are: - LocalLoadArgs - HFLoadArgs - ArweaveLoadArgs | required

Returns:

Name | Type | Description
--- | --- | ---
`str` | `str` | the path to the model

Source code in `src/infernet_ml/utils/model_loader.py`

    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    def download_model(
        model_source: ModelSource,
        load_args: LoadArgs,
    ) -> str:
        """
        Load the model from the specified source.

        Args:
            model_source (ModelSource): the source of the model
            load_args (LoadArgs): the load arguments, options are:
                - LocalLoadArgs
                - HFLoadArgs
                - ArweaveLoadArgs

        Returns:
            str: the path to the model
        """
        logger.info(f"Downloading model from {model_source} with args {load_args}")

        match model_source:
            # load the model locally
            case ModelSource.LOCAL:
                local_args = cast(LocalLoadArgs, load_args)
                return local_args.path
            case ModelSource.HUGGINGFACE_HUB:
                hf_args = cast(HFLoadArgs, load_args)
                return cast(
                    str,
                    hf_hub_download(
                        repo_id=hf_args.repo_id,
                        filename=hf_args.filename,
                        revision=hf_args.version,
                        cache_dir=hf_args.cache_path,
                    ),
                )
            case ModelSource.ARWEAVE:
                arweave_args = cast(ArweaveLoadArgs, load_args)
                cache_path = arweave_args.cache_path or os.path.expanduser("~/.cache/")
                version = arweave_args.version or "latest"
                base_path = f"{cache_path}/{arweave_args.repo_id}/{version}/."
                logging.info(
                    f"Downloading model from Arweave "
                    f"{cache_path}/{arweave_args.repo_id}/{version}/{arweave_args.filename}"
                )
                return RepoManager().download_artifact_file(
                    repo_id=arweave_args.repo_id,
                    file_name=arweave_args.filename,
                    version=arweave_args.version,
                    base_path=base_path,
                )
            case _:
                raise ValueError(f"Invalid model source {model_source}")

## `parse_load_args(model_source, config)`

Parse the load arguments based on the model source.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`model_source` | `ModelSource` | the source of the model | required
`config` | `dict[str, str]` | the configuration | required

Returns:

Name | Type | Description
--- | --- | ---
`LoadArgs` | `LoadArgs` | the load arguments

Raises:

Type | Description
--- | ---
`ValueError` | if the model source is invalid

Source code in `src/infernet_ml/utils/model_loader.py`

     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    def parse_load_args(model_source: ModelSource, config: Any) -> LoadArgs:
        """
        Parse the load arguments based on the model source.

        Args:
            model_source (ModelSource): the source of the model
            config (dict[str, str]): the configuration

        Returns:
            LoadArgs: the load arguments

        Raises:
            ValueError: if the model source is invalid
        """

        match model_source:
            # parse the load arguments for the local model
            case ModelSource.LOCAL:
                return LocalLoadArgs(path=config["model_path"])
            # parse the load arguments for the model from Hugging Face Hub
            case ModelSource.HUGGINGFACE_HUB:
                return HFLoadArgs(repo_id=config["repo_id"], filename=config["filename"])
            # parse the load arguments for the model from Arweave
            case ModelSource.ARWEAVE:
                return ArweaveLoadArgs(
                    repo_id=config["repo_id"],
                    filename=config["filename"],
                    version=config.get("version"),
                )
            case _:
                raise ValueError(f"Invalid model source {model_source}")