hubconf.py

""" Configuration for torch hub usage """

from typing import Callable, Dict, Any

import torch

dependencies = ['torch', 'trecover']


def add_available_versions_to_docstring(handler: Callable[..., Callable]) -> Callable:
    from trecover.config import var

    handler.__doc__ += '\n'.join([f'\t* {version}' for version in var.CHECKPOINT_URLS.keys()])

    return handler


@add_available_versions_to_docstring
def trecover(device: torch.device = torch.device('cpu'), version: str = 'latest'):
    """
    Load the TRecover model via torch.hub.

    Parameters
    ----------
    device : torch.device, default=torch.device('cpu')
        Device on which to allocate the model.
    version : str, default='latest'
        Model weights' version.

    Returns
    -------
    model : TRecover
        Model with specified weights' version.

    Examples
    --------
    Show the docstring with available versions for the TRecover model:
        >>> print(torch.hub.help(github='alex-snd/TRecover', model='trecover'))

    Load the TRecover model:
        >>> torch.hub.load('alex-snd/TRecover', model='trecover', device=torch.device('cpu'), version='latest')

    Available Versions
    ------------------
    """

    import json
    from urllib.request import urlopen
    from trecover.config import var
    from trecover.model import TRecover

    if version not in var.CHECKPOINT_URLS:
        version = 'latest'

    with urlopen(var.CHECKPOINT_URLS[version]['config']) as url:
        config = json.loads(url.read().decode())

    model = TRecover(token_size=config['token_size'], pe_max_len=config['pe_max_len'],
                     num_layers=config['num_layers'], d_model=config['d_model'], n_heads=config['n_heads'],
                     d_ff=config['d_ff'], dropout=config['dropout'])

    model.load_state_dict(torch.hub.load_state_dict_from_url(url=var.CHECKPOINT_URLS[version]['model'],
                                                             progress=False,
                                                             map_location=device))

    return model


def collab_args() -> Dict[str, Any]:
    """
    Load collaborative training base arguments.

    Returns
    -------
    Dict[str, Any] :
        Collaborative training base arguments.

    """

    return {
        'initial_peers': [
            '/ip4/95.216.202.215/tcp/34234/p2p/QmS2mYPX8Q78RDQzxMf17phVYPashi4C6ixVbA3jrj9yxt',
        ],
        'experiment_prefix': 'trecover',
        'target_batch_size': 4096,
        'min_noise': 0,
        'max_noise': 1,  # TODO model params

    }