"""Update the loadables dict from the given mapping.
Arguments
---------
loadables : mapping
Mapping from loadable key to object
"""
self.loadables.update(loadables)
defadd_paths(self, paths):
"""Update the paths for different loadables.
When collecting parameters, paths here are preferred. Note that when
collecting, you can specify a default source, which is used for all
loadables that don't have a path specified.
Arguments
---------
paths : mapping
Mapping from loadable key to filepath. The last part
of the path is treated as file name, the rest of it
is treated as a "source" which can be either a directory
path or a magic source like Huggingface hub ID.
e.g. sb/asr-crdnn-libri/lm.ckpt
-> source=sb/asr-crdnn-libri, file=lm.ckpt
"""
self.paths.update(paths)
defadd_custom_hooks(self, custom_hooks):
"""Update the custom hooks.
When loading parameters, hooks here are preferred over class defaults.
Arguments
---------
custom_hooks : mapping
Mapping from loadable key to parameter transfer hook function. If
you want to use a custom loading function, specify it here.
"""
self.custom_hooks.update(custom_hooks)
defadd_conditions(self, conditions):
"""Update the conditions.
Arguments
---------
conditions: mapping
Mapping from loadable keys to condition values,
useful for loading certain elements only if a flag is turned on
"""
self.conditions.update(conditions)
@staticmethod
defsplit_path(path):
"""Splits a path to source and filename
This also handles URLs and Huggingface hub paths, in addition to
regular paths.
Arguments
---------
path : str
Returns
-------
str
Source
str
Filename
"""
defsplit(src):
"""Core function to split path."""
if"/"insrc:
returnsrc.rsplit("/", maxsplit=1)
else:
# Interpret as path to file in current directory.
return"./", src
ifisinstance(path, FetchSource):
fetch_from, fetch_path=path
source, filename=split(fetch_path)
returnFetchSource(fetch_from, source), filename
else:
returnsplit(path)
defcollect_files(
self,
default_source=None,
local_strategy=LocalStrategy.SYMLINK,
fetch_config=FetchConfig(),
):
"""Fetches parameters from known paths with fallback default_source
The actual parameter files may reside elsewhere, but this ensures a
symlink in the self.collect_in directory. The symlink always uses the
loadable key in the filename. This standardization makes it easier to
orchestrate pretraining on e.g. distributed setups.
Use the default_source if you have everything organized neatly into one
location, like a Huggingface hub repo.
Arguments
---------
default_source : str or Path or FetchSource
This is used for each loadable which doesn't have a path already
specified.
e.g. if the loadable has key `"asr"`, then the file to look for is
`<default_source>/asr.ckpt`
local_strategy : LocalStrategy
How to perform caching on the file for local storage.
fetch_config : FetchConfig
Configuration options like caching strategy for fetching files.
Returns
-------
dict
Mapping from loadable key to a local path from which loadable's
parameters can be loaded. This is not used in this class, but
can possibly be helpful.
"""
ifself.collect_inisnotNone:
logger.debug(
f"Collecting files (or symlinks) for pretraining in {self.collect_in}."
)
self.collect_in.mkdir(exist_ok=True)
if (
platform.system() =="Windows"
andlocal_strategy==LocalStrategy.SYMLINK
):
warnings.warn(
"Requested Pretrainer collection using symlinks on Windows. This might not work; see `LocalStrategy` documentation. Consider unsetting `collect_in` in Pretrainer to avoid symlinking altogether."
)
else:
logger.debug(
"Fetching files for pretraining (no collection directory set)"
f"Redirecting (loading from local path): {name} -> {self.paths[name]}"
)
paramfiles[name] =self.paths[name]
elifself.collect_inisnotNone:
paramfiles[name] =self.collect_in/filename
else:
raiseValueError(
f'Pretrainer has never collected `{name}`, did you forget a call to `collect_files`? Could not fall back to `collect_in`, as it was not specified (default is no longer "model_checkpoints").'
)
self._call_load_hooks(paramfiles)
def_call_load_hooks(self, paramfiles):
# This internal function finds the correct hook to call for every