Function load

src/models/text_decoder_v2.py:494–518

Loads pre-trained model weights from a checkpoint file. Args: init_params: A PyTree of randomly initialized model parameters. init_file: Path to the checkpoint file to load. model_cfg: The model configuration dictionary (unused in this function). dont_load: A seq

(init_params: dict, init_file: str, model_cfg: dict, dont_load: Sequence[str] = ())

def load(init_params: dict, init_file: str, model_cfg: dict, dont_load: Sequence[str] = ()) -> dict:
    """Loads pre-trained model weights from a checkpoint file.

    Args:
        init_params: A PyTree of randomly initialized model parameters.
        init_file: Path to the checkpoint file to load.
        model_cfg: The model configuration dictionary (unused in this function).
        dont_load: A sequence of parameter names (or prefixes) to exclude from loading.

    Returns:
        A PyTree of parameters with loaded weights.
    """
    del model_cfg
    restored_params = utils.load_params(filepath=init_file)

    # Merge restored params into the initialized structure, allowing for fine-tuning
    # where some parameters (like the output head) might be re-initialized.
    restored_params = common.merge_params(
        restored_params, init_params, dont_load=dont_load)

    # Ensure dtypes are correctly restored (e.g., from float32 to bfloat16).
    restored_params = jax.tree_util.tree_map(
        utils.recover_dtype, restored_params)
        
    return restored_params