# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.## Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License."""Interfaces common to all Neural Modules and Models."""from__future__importannotationsimportcopyimporthashlibimportinspectimportosimportshutilimporttracebackfromabcimportABC,abstractmethodfromcollections.abcimportMapping,Sequencefromcontextlibimportcontextmanagerfromdataclassesimportdataclass,field,is_dataclassfromenumimportEnumfromfunctoolsimporttotal_orderingfrompathlibimportPath,PurePosixPathfromtypingimportAny,Dict,List,Optional,Tuple,Unionimporthydraimporttorchimportwraptfromhuggingface_hubimport_CACHED_NO_EXIST,HfApifromhuggingface_hubimportget_tokenasget_hf_tokenfromhuggingface_hubimporthf_hub_download,snapshot_download,try_to_load_from_cachefromomegaconfimportDictConfig,OmegaConfimportnemofromnemo.core.classes.mixins.hf_io_mixinimportHuggingFaceFileIOfromnemo.core.config.templates.model_cardimportNEMO_DEFAULT_MODEL_CARD_TEMPLATEfromnemo.core.connectors.save_restore_connectorimportSaveRestoreConnectorfromnemo.core.neural_typesimportNeuralType,NeuralTypeComparisonResultfromnemo.utilsimportlogging,model_utilsfromnemo.utils.cloudimportmaybe_download_from_cloudfromnemo.utils.data_utilsimportresolve_cache_dirfromnemo.utils.model_utilsimportimport_class_by_path,maybe_update_config_version__all__=['Typing','FileIO','Model','Serialization','typecheck','PretrainedModelInfo']_TYPECHECK_ENABLED=True_TYPECHECK_SEMANTIC_CHECK_ENABLED=True# Added these for now but these should be updated based on collectionsALLOWED_TARGET_PREFIXES=["nemo.collections.","nemo.core.","nemo.utils.","nemo.lightning.","nemo_text_processing.text_normalization.normalize.Normalizer","tests.collections.","tests.core.","torch.nn.","torch.distributed.fsdp.","torch.optim.","torch.utils.data.","torchmetrics.","lightning.pytorch.callbacks.","lightning.pytorch.loggers.","lightning.pytorch.strategies.","lightning.pytorch.accelerators.","omegaconf.","megatron.",]ALLOWED_CALLABLE_PREFIXES=["nemo.collections.common.tokenizers","nemo.collections.common.parts","nemo.collections.asr.modules","nemo.collections.asr.parts","nemo.collections.audio.parts","nemo.collections.speechlm","nemo.collections.llm","nemo.lightning","megatron.core","tests.collections.llm.common",]ALLOWED_ADAPTER_STRATEGY_PREFIXES=["nemo.core.classes.mixins.adapter_mixin_strategies","nemo.collections.asr.parts.submodules.adapters",]ALLOWED_CLASS_PREFIXES_WITH_OPTIONAL_DEPENDENCIES=["nemo.collections.audio.parts.submodules.flow","nemo.collections.common.tokenizers","nemo.collections.speechlm2.parts.parallel","nemo.collections.tts.g2p",]ALLOWED_EXACT_CLASS_TARGETS={"nemo_text_processing.text_normalization.normalize.Normalizer",}ALLOWED_LEGACY_FALLBACK_TARGETS={"src.multi_classification_models.EncDecMultiClassificationModel",}classUnsafeTargetError(ValueError):"""Raised when config-driven instantiation requests a disallowed target."""def_is_target_allowed(target:str)->bool:""" Return True if the Hydra `_target_` should be allowed to be instantiated. """# cheap prefix checkifnotany(target.startswith(prefix)forprefixinALLOWED_TARGET_PREFIXES):returnFalse# resolve to objecttry:obj=hydra.utils.get_class(target)exceptException:# Hydra fails on functions; try get_object insteadtry:obj=hydra.utils.get_object(target)exceptExceptionase2:# For NeMo targets that passed prefix check, be more lenient with import errors# This handles cases where dependencies might be missing during testingiftarget.startswith("nemo."):# Check if this is a missing dependency issue vs a malicious targeterror_msg=str(e2).lower()ifany(missing_depinerror_msgformissing_depin['no module named','modulenotfounderror']):# This appears to be a legitimate NeMo target with missing dependencies# Apply additional checks based on the target path structuretarget_parts=target.split('.')iflen(target_parts)>=3:# e.g., nemo.collections.asrmodule_path='.'.join(target_parts[:-1])# Remove function/class name# Check if the module path is in one of our approved prefixes.if(any(module_path.startswith(p)forpinALLOWED_CALLABLE_PREFIXES)orany(module_path.startswith(p)forpinALLOWED_ADAPTER_STRATEGY_PREFIXES)orany(module_path.startswith(p)forpinALLOWED_CLASS_PREFIXES_WITH_OPTIONAL_DEPENDENCIES)):# This is likely a legitimate NeMo function/class that we can't import# due to missing dependencies. We'll assume it's safe.returnTruereturnFalse# @experimental / @deprecated wrap the class in a wrapt proxy that passes# isinstance(.., type) but breaks issubclass(); unwrap to the real class.whilehasattr(obj,"__wrapped__"):obj=obj.__wrapped__# If it's a class: allow only subclasses of safe basesifisinstance(obj,type):iftarget.startswith("nemo.core.config.")andis_dataclass(obj):returnTruefromnemo.core.classes.modelPTimportModelPTiftargetinALLOWED_EXACT_CLASS_TARGETS:returnTrueserialization_cls=globals().get("Serialization")ifserialization_clsisnotNone:try:ifissubclass(obj,serialization_cls):returnTrueexceptTypeError:returnFalseSAFE_BASES=(torch.nn.Module,ModelPT)try:ifissubclass(obj,SAFE_BASES):returnTrueexceptTypeError:returnFalsetry:ifissubclass(obj,torch.utils.data.Dataset):returnTrueexceptTypeError:returnFalseiftarget.startswith("torch.optim."):try:returnissubclass(obj,torch.optim.Optimizer)exceptTypeError:returnFalseiftarget.startswith("torchmetrics."):try:fromtorchmetricsimportMetricreturnissubclass(obj,Metric)except(ImportError,TypeError):returnFalseiftarget=="torch.distributed.fsdp.MixedPrecisionPolicy":returnis_dataclass(obj)module_name=getattr(obj,"__module__","")or""ifany(module_name.startswith(p)forpinALLOWED_ADAPTER_STRATEGY_PREFIXES):fromnemo.core.classes.mixins.adapter_mixin_strategiesimportAbstractAdapterStrategytry:returnissubclass(obj,AbstractAdapterStrategy)exceptTypeError:returnFalseiftarget.startswith("nemo.collections.common.tokenizers.")ortarget.startswith("nemo.collections.tts.torch.tts_tokenizers."):try:fromnemo.collections.common.tokenizers.text_to_speech.tts_tokenizersimportBaseTokenizerfromnemo.collections.common.tokenizers.tokenizer_specimportTokenizerSpecreturnissubclass(obj,(BaseTokenizer,TokenizerSpec))except(ImportError,TypeError):returnFalseiftarget.startswith("nemo.collections.tts.g2p.")ortarget=="nemo.collections.tts.torch.g2ps.EnglishG2p":try:fromnemo.collections.tts.g2p.models.baseimportBaseG2preturnissubclass(obj,BaseG2p)except(ImportError,TypeError):returnFalseiftarget.startswith("nemo.collections.tts.parts.preprocessing."):try:fromnemo.collections.tts.parts.preprocessing.audio_trimmingimportAudioTrimmerreturnissubclass(obj,AudioTrimmer)except(ImportError,TypeError):returnFalseiftarget.startswith("nemo.collections.tts.parts.utils.callbacks"):try:fromnemo.collections.tts.parts.utils.callbacksimportArtifactGeneratorreturnissubclass(obj,ArtifactGenerator)except(ImportError,TypeError):returnFalseiftarget.startswith("nemo.collections.audio.parts.submodules.flow."):try:fromnemo.collections.audio.parts.submodules.flowimport(ConditionalFlow,ConditionalFlowMatchingSampler,)returnissubclass(obj,(ConditionalFlow,ConditionalFlowMatchingSampler))except(ImportError,TypeError):returnFalseiftarget.startswith("nemo.core.optim.lr_scheduler."):try:fromtorch.optim.lr_schedulerimport_LRSchedulerreturnissubclass(obj,_LRScheduler)except(ImportError,TypeError):returnFalseiftarget.startswith("nemo.collections.speechlm2.parts.parallel."):try:fromlightning.pytorch.strategies.model_parallelimportModelParallelStrategyreturnissubclass(obj,ModelParallelStrategy)except(ImportError,TypeError):returnFalseiftarget.startswith("lightning.pytorch."):try:iftarget.startswith("lightning.pytorch.accelerators."):fromlightning.pytorch.acceleratorsimportAcceleratorreturnissubclass(obj,Accelerator)iftarget.startswith("lightning.pytorch.callbacks."):fromlightning.pytorch.callbacksimportCallbackreturnissubclass(obj,Callback)iftarget.startswith("lightning.pytorch.loggers."):fromlightning.pytorch.loggers.loggerimportLoggerreturnissubclass(obj,Logger)iftarget.startswith("lightning.pytorch.strategies."):fromlightning.pytorch.strategiesimportStrategyreturnissubclass(obj,Strategy)except(ImportError,TypeError):returnFalse# If it's a callable function: allow only if in approved submodules.ifcallable(obj)andnotisinstance(obj,type):module_name=getattr(obj,"__module__","")or""ifany(module_name.startswith(p)forpinALLOWED_CALLABLE_PREFIXES):returnTruereturnFalse# otherwise disallowreturnFalsedef_unsafe_target_error(target_path:str,config_key:str)->ValueError:returnUnsafeTargetError(f"Instantiation of unsafe target '{target_path}' is blocked. "f"The '{config_key}' must point to a class or function within an approved namespace. "f"This restriction is in place to prevent potential arbitrary code execution.")def_get_allowed_target_class(target_path:str):ifnot_is_target_allowed(target_path):raise_unsafe_target_error(target_path,"target")returnimport_class_by_path(target_path)def_validate_config_targets_recursive(config_node:Any):ifisinstance(config_node,Mapping):# Handles DictConfig and dictif"_target_"inconfig_node:target_path=config_node["_target_"]ifnot_is_target_allowed(target_path):raise_unsafe_target_error(target_path,"_target_")forkey,valueinconfig_node.items():_validate_config_targets_recursive(value)elifisinstance(config_node,Sequence)andnotisinstance(config_node,str):# Handles ListConfig and listforiteminconfig_node:_validate_config_targets_recursive(item)defsafe_instantiate(config:DictConfig,*args,**kwargs):""" A wrapper around hydra.utils.instantiate that first validates all _target_ fields in the config against an allow-list of prefixes. """ifconfigisnotNone:_validate_config_targets_recursive(config)returnhydra.utils.instantiate(config,*args,**kwargs)defis_typecheck_enabled():""" Getter method for typechecking state. """return_TYPECHECK_ENABLEDdefis_semantic_typecheck_enabled():""" Getter method for typechecking semantics state. """return_TYPECHECK_SEMANTIC_CHECK_ENABLED@dataclassclassTypecheckMetadata:""" Metadata class for input/output neural types. # Primary attributes original_types: Preserve the dictionary of type information provided. ignore_collections: For backward compatibility, container support can be disabled explicitly using this flag. When set to True, all nesting is ignored and nest-depth checks are skipped. # Derived attributed mandatory_types: Sub-dictionary of `original_types` which contains only those types which are mandatory to include when calling the function. base_types: Dictionary of flattened `str: NeuralType` definitions, disregarding the nest level details into appropriate arguments. container_depth: Dictionary mapping `str: int` - such that the valid depth of the nest of this neural type is recorded. has_container_types: Bool flag declaring if any of the neural types declares a container nest in its signature. is_singular_container_type: Bool flag declaring if this is a single Neural Type with a container nest in its signature. Required for supporting python list expansion in return statement. """original_types:Dict[str,NeuralType]ignore_collections:boolmandatory_types:Dict[str,NeuralType]=field(init=False)base_types:Dict[str,NeuralType]=field(init=False)container_depth:Dict[str,int]=field(init=False)has_container_types:bool=field(init=False)is_singular_container_type:bool=field(init=False)def__post_init__(self):# If even one NeuralType declares a container nest, set to Truehas_container_types=Falsefortype_valinself.original_types.values():ifisinstance(type_val,(list,tuple)):has_container_types=Truebreakself.has_container_types=has_container_types# If only one NeuralType is declared, and it declares a container nest, set to Trueifself.has_container_typesandlen(self.original_types)==1:self.is_singular_container_type=Trueelse:self.is_singular_container_type=False# If container nests are declared, flatten the nest into `base_types`# Also compute the nest depth for each of the NeuralTypesifself.has_container_types:self.base_types={}self.container_depth={}fortype_key,type_valinself.original_types.items():depth=0whileisinstance(type_val,(list,tuple)):iflen(type_val)>1:raiseTypeError(f"Neural Type `{type_key}`: {type_val} definition contains more than one element when ""declaring the nested container structure.\n""Please ensure that you have only 1 NeuralType inside of the entire nested structure ""definition.")type_val=type_val[0]depth+=1self.base_types[type_key]=type_valself.container_depth[type_key]=depthelse:# Otherwise, simply preserve the original_types and set depth of nest to 0.self.base_types=self.original_typesself.container_depth={type_key:0fortype_keyinself.base_types.keys()}# Compute subset of original_types which are mandatory in the call argspecself.mandatory_types={type_key:type_valfortype_key,type_valinself.base_types.items()ifnottype_val.optional}
classTyping(ABC):""" An interface which endows module with neural types """@propertydefinput_types(self)->Optional[Dict[str,NeuralType]]:"""Define these to enable input neural type checks"""returnNone@propertydefoutput_types(self)->Optional[Dict[str,NeuralType]]:"""Define these to enable output neural type checks"""returnNone
def_validate_input_types(self,input_types=None,ignore_collections=False,**kwargs):""" This function does a few things. 1) It ensures that len(self.input_types <non-optional>) <= len(kwargs) <= len(self.input_types). 2) For each (keyword name, keyword value) passed as input to the wrapped function: - Check if the keyword name exists in the list of valid self.input_types names. - Check if keyword value has the `neural_type` property. - If it does, then perform a comparative check and assert that neural types are compatible (SAME or GREATER). - Check if keyword value is a container type (list or tuple). If yes, then perform the elementwise test of neural type above on each element of the nested structure, recursively. Args: input_types: Either the `input_types` defined at class level, or the local function overridden type definition. ignore_collections: For backward compatibility, container support can be disabled explicitly using this flag. When set to True, all nesting is ignored and nest-depth checks are skipped. kwargs: Dictionary of argument_name:argument_value pairs passed to the wrapped function upon call. """ifinput_typesisnotNone:# Precompute metadatametadata=TypecheckMetadata(original_types=input_types,ignore_collections=ignore_collections)total_input_types=len(input_types)mandatory_input_types=len(metadata.mandatory_types)# Allow number of input arguments to be <= total input neural types.iflen(kwargs)<mandatory_input_typesorlen(kwargs)>total_input_types:raiseTypeError(f"Number of input arguments provided ({len(kwargs)}) is not as expected. Function has "f"{total_input_types} total inputs with {mandatory_input_types} mandatory inputs.")forkey,valueinkwargs.items():# Check if keys exists in the defined input typesifkeynotininput_types:raiseTypeError(f"Input argument {key} has no corresponding input_type match. "f"Existing input_types = {input_types.keys()}")# Perform neural type checkif(hasattr(value,'neural_type')andis_semantic_typecheck_enabled()andnotmetadata.base_types[key].compare(value.neural_type)in(NeuralTypeComparisonResult.SAME,NeuralTypeComparisonResult.GREATER,)):error_msg=[f"{input_types[key].compare(value.neural_type)} :",f"Input type expected : {input_types[key]}",f"Input type found : {value.neural_type}",f"Argument: {key}",]fori,dict_tupleinenumerate(metadata.base_types[key].elements_type.type_parameters.items()):error_msg.insert(i+2,f' input param_{i} : {dict_tuple[0]}: {dict_tuple[1]}')fori,dict_tupleinenumerate(value.neural_type.elements_type.type_parameters.items()):error_msg.append(f' input param_{i} : {dict_tuple[0]}: {dict_tuple[1]}')raiseTypeError("\n".join(error_msg))# Perform input ndim checkifhasattr(value,'shape'):value_shape=value.shapetype_shape=metadata.base_types[key].axesname=keyiftype_shapeisnotNoneandlen(value_shape)!=len(type_shape):raiseTypeError(f"Input shape mismatch occured for {name} in module {self.__class__.__name__} : \n"f"Input shape expected = {metadata.base_types[key].axes} | \n"f"Input shape found : {value_shape}")# Perform recursive neural type check for homogeneous elementselifisinstance(value,list)orisinstance(value,tuple):forind,valinenumerate(value):""" This initiates a DFS, tracking the depth count as it goes along the nested structure. Initial depth is 1 as we consider the current loop to be the 1st step inside the nest. """self.__check_neural_type(val,metadata,depth=1,name=key)
def_attach_and_validate_output_types(self,out_objects,ignore_collections=False,output_types=None):""" This function does a few things. 1) It ensures that len(out_object) == len(self.output_types). 2) If the output is a tensor (or list/tuple of list/tuple ... of tensors), it attaches a neural_type to it. For objects without the neural_type attribute, such as python objects (dictionaries and lists, primitive data types, structs), no neural_type is attached. Note: tensor.neural_type is only checked during _validate_input_types which is called prior to forward(). Args: output_types: Either the `output_types` defined at class level, or the local function overridden type definition. ignore_collections: For backward compatibility, container support can be disabled explicitly using this flag. When set to True, all nesting is ignored and nest-depth checks are skipped. out_objects: The outputs of the wrapped function. """# TODO: Properly implement thisifoutput_typesisnotNone:# Precompute metadatametadata=TypecheckMetadata(original_types=output_types,ignore_collections=ignore_collections)out_types_list=list(metadata.base_types.items())mandatory_out_types_list=list(metadata.mandatory_types.items())# First convert all outputs to list/tuple format to check correct number of outputsifisinstance(out_objects,(list,tuple)):out_container=out_objects# can be any rank nested structureelse:out_container=[out_objects]# If this neural type has a *single output*, with *support for nested outputs*,# then *do not* perform any check on the number of output items against the number# of neural types (in this case, 1).# This is done as python will *not* wrap a single returned list into a tuple of length 1,# instead opting to keep the list intact. Therefore len(out_container) in such a case# is the length of all the elements of that list - each of which has the same corresponding# neural type (defined as the singular container type).ifmetadata.is_singular_container_type:pass# In all other cases, python will wrap multiple outputs into an outer tuple.# Allow number of output arguments to be <= total output neural types and >= mandatory outputs.eliflen(out_container)>len(out_types_list)orlen(out_container)<len(mandatory_out_types_list):raiseTypeError("Number of output arguments provided ({}) is not as expected. ""It should be larger or equal than {} and less or equal than {}.\n""This can be either because insufficient/extra number of output NeuralTypes were provided,""or the provided NeuralTypes {} should enable container support ""(add '[]' to the NeuralType definition)".format(len(out_container),len(out_types_list),len(mandatory_out_types_list),output_types))# Attach types recursively, if possibleifnotisinstance(out_objects,tuple)andnotisinstance(out_objects,list):# Here, out_objects is a single object which can potentially be attached with a NeuralTypetry:out_objects.neural_type=out_types_list[0][1]exceptException:pass# Perform output ndim checkifhasattr(out_objects,'shape'):value_shape=out_objects.shapetype_shape=out_types_list[0][1].axesname=out_types_list[0][0]iftype_shapeisnotNoneandlen(value_shape)!=len(type_shape):raiseTypeError(f"Output shape mismatch occured for {name} in module {self.__class__.__name__} : \n"f"Output shape expected = {type_shape} | \n"f"Output shape found : {value_shape}")elifmetadata.is_singular_container_type:# If only a single neural type is provided, and it defines a container nest,# then all elements of the returned list/tuple are assumed to belong to that# singular neural type.# As such, the "current" depth inside the DFS loop is counted as 1,# and subsequent nesting will increase this count.# NOTE:# As the flag `is_singular_container_type` will activate only for# the case where there is 1 output type defined with container nesting,# this is a safe assumption to make.depth=1# NOTE:# A user may chose to explicitly wrap the single output list within an explicit tuple# In such a case we reduce the "current" depth to 0 - to acknowledge the fact that# the actual nest exists within a wrapper tuple.iflen(out_objects)==1andtype(out_objects)==tuple:depth=0forind,resinenumerate(out_objects):self.__attach_neural_type(res,metadata,depth=depth,name=out_types_list[0][0])else:# If more then one item is returned in a return statement, python will wrap# the output with an outer tuple. Therefore there must be a 1:1 correspondence# of the output_neural type (with or without nested structure) to the actual output# (whether it is a single object or a nested structure of objects).# Therefore in such a case, we "start" the DFS at depth 0 - since the recursion is# being applied on 1 neural type : 1 output struct (single or nested output).# Since we are guarenteed that the outer tuple will be built by python,# assuming initial depth of 0 is appropriate.forind,resinenumerate(out_objects):self.__attach_neural_type(res,metadata,depth=0,name=out_types_list[ind][0])
def__check_neural_type(self,obj,metadata:TypecheckMetadata,depth:int,name:str=None):""" Recursively tests whether the obj satisfies the semantic neural type assertion. Can include shape checks if shape information is provided. Args: obj: Any python object that can be assigned a value. metadata: TypecheckMetadata object. depth: Current depth of recursion. name: Optional name used of the source obj, used when an error occurs. """ifisinstance(obj,tuple)orisinstance(obj,list):foreleminobj:self.__check_neural_type(elem,metadata,depth+1,name=name)return# after processing nest, return to avoid testing nest itselftype_val=metadata.base_types[name]# If nest depth doesnt match neural type structure depth, raise an errorifnotmetadata.ignore_collectionsanddepth!=metadata.container_depth[name]:raiseTypeError("While checking input neural types,\n""Nested depth of value did not match container specification:\n"f"Current nested depth of NeuralType '{name}' ({type_val}): {depth}\n"f"Expected nested depth : {metadata.container_depth[name]}")if(hasattr(obj,'neural_type')andis_semantic_typecheck_enabled()andnottype_val.compare(obj.neural_type)in(NeuralTypeComparisonResult.SAME,NeuralTypeComparisonResult.GREATER,)):raiseTypeError(f"{type_val.compare(obj.neural_type)} : \n"f"Input type expected = {type_val} | \n"f"Input type found : {obj.neural_type}")# Perform input ndim checkifhasattr(obj,'shape'):value_shape=obj.shapetype_shape=type_val.axesiftype_shapeisnotNoneandlen(value_shape)!=len(type_shape):raiseTypeError(f"Input shape mismatch occured for {name} in module {self.__class__.__name__} : \n"f"Input shape expected = {type_shape} | \n"f"Input shape found : {value_shape}")def__attach_neural_type(self,obj,metadata:TypecheckMetadata,depth:int,name:str=None):""" Recursively attach neural types to a given object - as long as it can be assigned some value. Args: obj: Any python object that can be assigned a value. metadata: TypecheckMetadata object. depth: Current depth of recursion. name: Optional name used of the source obj, used when an error occurs. """ifisinstance(obj,tuple)orisinstance(obj,list):foreleminobj:self.__attach_neural_type(elem,metadata,depth=depth+1,name=name)return# after processing nest, return to avoid argument insertion into nest itselftype_val=metadata.base_types[name]# If nest depth doesnt match neural type structure depth, raise an errorifnotmetadata.ignore_collectionsanddepth!=metadata.container_depth[name]:raiseTypeError("While attaching output neural types,\n""Nested depth of value did not match container specification:\n"f"Current nested depth of NeuralType '{name}' ({type_val}): {depth}\n"f"Expected nested depth : {metadata.container_depth[name]}")try:obj.neural_type=type_valexceptException:pass# Perform output ndim checkifhasattr(obj,'shape'):value_shape=obj.shapetype_shape=type_val.axesiftype_shapeisnotNoneandlen(value_shape)!=len(type_shape):raiseTypeError(f"Output shape mismatch occured for {name} in module {self.__class__.__name__} : \n"f"Output shape expected = {type_shape} | \n"f"Output shape found : {value_shape}")
classSerialization(ABC):# pylint: disable=C0115
@classmethoddeffrom_config_dict(cls,config:'DictConfig',trainer:Optional['Trainer']=None):# noqa: F821"""Instantiates object using DictConfig-based configuration"""# Resolve the config dictifisinstance(config,DictConfig):config=model_utils.convert_model_config_to_dict_config(config)config=maybe_update_config_version(config,make_copy=False)# Hydra 0.x APIif('cls'inconfigor'target'inconfig)and'params'inconfig:# regular hydra-based instantiationinstance=safe_instantiate(config=config)# Hydra 1.x APIelif'_target_'inconfig:# regular hydra-based instantiationinstance=safe_instantiate(config=config)else:instance=Noneprev_error=""# Attempt class path resolution from config `target` class (if it exists)if'target'inconfig:target_cls_path=config["target"]# No guarantee that this is a omegaconf classimported_cls=Nonetry:iftarget_cls_pathinALLOWED_LEGACY_FALLBACK_TARGETS:imported_cls=clselse:# try to import the target classimported_cls=_get_allowed_target_class(target_cls_path)# if calling class (cls) is subclass of imported class,# use subclass insteadifissubclass(cls,imported_cls):imported_cls=clsaccepts_trainer=Serialization._inspect_signature_for_trainer(imported_cls)ifaccepts_trainer:instance=imported_cls(cfg=config,trainer=trainer)else:instance=imported_cls(cfg=config)exceptUnsafeTargetError:raiseexceptExceptionase:# record previous errortb=traceback.format_exc()prev_error=(f"Model instantiation failed!\nTarget class:\t{target_cls_path}"f"\nError(s):\t{e}\n{tb}")logging.debug(prev_error+"\nFalling back to `cls`.")# target class resolution was unsuccessful, fall back to current `cls`ifinstanceisNone:try:accepts_trainer=Serialization._inspect_signature_for_trainer(cls)ifaccepts_trainer:instance=cls(cfg=config,trainer=trainer)else:instance=cls(cfg=config)exceptExceptionase:# report saved errors, if any, and raiseifprev_error:logging.error(prev_error)raiseeifnothasattr(instance,'_cfg'):instance._cfg=configreturninstance
defto_config_dict(self)->'DictConfig':"""Returns object's configuration to config dictionary"""ifhasattr(self,'_cfg')andself._cfgisnotNone:# Resolve the config dictconfig=model_utils.convert_model_config_to_dict_config(self._cfg)config=maybe_update_config_version(config,make_copy=False)self._cfg=configreturnself._cfgelse:raiseNotImplementedError('to_config_dict() can currently only return object._cfg but current object does not have it.')
@classmethoddef_inspect_signature_for_trainer(cls,check_cls):ifhasattr(check_cls,'__init__'):signature=inspect.signature(check_cls.__init__)if'trainer'insignature.parameters:returnTrueelse:returnFalseelse:returnFalse
classFileIO(ABC):# pylint: disable=C0115
defsave_to(self,save_path:str):""" Standardized method to save a tarfile containing the checkpoint, config, and any additional artifacts. Implemented via :meth:`nemo.core.connectors.save_restore_connector.SaveRestoreConnector.save_to`. Args: save_path: str, path to where the file should be saved. """raiseNotImplementedError()
@classmethoddefrestore_from(cls,restore_path:str,override_config_path:Optional[str]=None,map_location:Optional['torch.device']=None,strict:bool=True,return_config:bool=False,trainer:Optional['Trainer']=None,# noqa: F821save_restore_connector:SaveRestoreConnector=None,):""" Restores model instance (weights and configuration) from a .nemo file Args: restore_path: path to .nemo file from which model should be instantiated override_config_path: path to a yaml config that will override the internal config file or an OmegaConf / DictConfig object representing the model config. map_location: Optional torch.device() to map the instantiated model to a device. By default (None), it will select a GPU if available, falling back to CPU otherwise. strict: Passed to load_state_dict. By default True return_config: If set to true, will return just the underlying config of the restored model as an OmegaConf DictConfig object without instantiating the model. trainer: An optional Trainer object, passed to the model constructor. save_restore_connector: An optional SaveRestoreConnector object that defines the implementation of the restore_from() method. """raiseNotImplementedError()
@classmethoddeffrom_config_file(cls,path2yaml_file:str):""" Instantiates an instance of NeMo Model from YAML config file. Weights will be initialized randomly. Args: path2yaml_file: path to yaml file with model configuration Returns: """ifissubclass(cls,Serialization):conf=OmegaConf.load(path2yaml_file)returncls.from_config_dict(config=conf)else:raiseNotImplementedError()
defto_config_file(self,path2yaml_file:str):""" Saves current instance's configuration to YAML config file. Weights will not be saved. Args: path2yaml_file: path2yaml_file: path to yaml file where model model configuration will be saved Returns: """ifhasattr(self,'_cfg'):self._cfg=maybe_update_config_version(self._cfg,make_copy=False)withopen(path2yaml_file,'w',encoding='utf-8')asfout:OmegaConf.save(config=self._cfg,f=fout,resolve=True)else:raiseNotImplementedError()
@total_ordering@dataclassclassPretrainedModelInfo:# pylint: disable=C0115pretrained_model_name:strdescription:strlocation:strclass_:'Model'=Nonealiases:List[str]=Nonedef__repr__(self):base=self.__class__.__name__extras=("pretrained_model_name={pretrained_model_name},\n\t""description={description},\n\t""location={location}".format(**self.__dict__))ifself.class_isnotNone:extras="{extras},\n\t""class_={class_}".format(extras=extras,**self.__dict__)representation=f"{base}(\n\t{extras}\n)"returnrepresentationdef__hash__(self):# assumes that locations are unique urls, and therefore their hashes# should ideally also be uniquelocation_hash=hash(self.location)returnlocation_hashdef__eq__(self,other):# another object is equal to self, iff# if it's hash is equal to hash(self)returnhash(self)==hash(other)orself.pretrained_model_name==other.pretrained_model_namedef__lt__(self,other):returnself.pretrained_model_name<other.pretrained_model_nameclassModel(Typing,Serialization,FileIO,HuggingFaceFileIO):""" Abstract class offering interface which should be implemented by all NeMo models. """@classmethod@abstractmethoddeflist_available_models(cls)->Optional[List[PretrainedModelInfo]]:""" Should list all pre-trained models available via NVIDIA NGC cloud. Note: There is no check that requires model names and aliases to be unique. In the case of a collision, whatever model (or alias) is listed first in the this returned list will be instantiated. Returns: A list of PretrainedModelInfo entries """pass@classmethoddefget_available_model_names(cls)->List[str]:""" Returns the list of model names available via NVIDIA NGC cloud, to get the complete model description use list_available_models() Returns: A list of model names """model_names=[]ifcls.list_available_models()isnotNone:model_names=[model.pretrained_model_nameformodelincls.list_available_models()]returnmodel_names@classmethoddeffrom_pretrained(cls,model_name:str,refresh_cache:bool=False,override_config_path:Optional[str]=None,map_location:Optional['torch.device']=None,strict:bool=True,return_config:bool=False,trainer:Optional['Trainer']=None,# noqa: F821save_restore_connector:SaveRestoreConnector=None,return_model_file:Optional[bool]=False,):""" Instantiates an instance of NeMo from NVIDIA NGC cloud Use restore_from() to instantiate from a local .nemo file. Args: model_name: string key which will be used to find the module. refresh_cache: If set to True, then when fetching from cloud, this will re-fetch the file from cloud even if it is already found in a cache locally. override_config_path: path to a yaml config that will override the internal config file map_location: Optional torch.device() to map the instantiated model to a device. By default (None), it will select a GPU if available, falling back to CPU otherwise. strict: Passed to torch.load_state_dict. By default true. return_config: If set to true, will return just the underlying config of the restored model as an OmegaConf DictConfig object without instantiating the model. return_model_file: If set to true, will return just the downloaded model file in cache Returns: A model instance of a particular model class or its underlying config (if return_config is set). """ifsave_restore_connectorisNone:save_restore_connector=SaveRestoreConnector()# Resolve if the pretrained model name is from NGC or other sources# HF Hub sourceif'/'inmodel_name:class_,nemo_model_file_in_cache=cls._get_hf_hub_pretrained_model_info(model_name=model_name,refresh_cache=refresh_cache)# Check if nemo_model_file_in_cache is a directoryifos.path.isdir(nemo_model_file_in_cache):# Update SaveRestoreConnector with the flag to read from an unpacked NeMo foldersave_restore_connector.model_extracted_dir=nemo_model_file_in_cacheelse:# NGC sourceclass_,nemo_model_file_in_cache=cls._get_ngc_pretrained_model_info(model_name=model_name,refresh_cache=refresh_cache)ifreturn_model_file:returnnemo_model_file_in_cacheinstance=class_.restore_from(restore_path=nemo_model_file_in_cache,override_config_path=override_config_path,map_location=map_location,strict=strict,return_config=return_config,trainer=trainer,save_restore_connector=save_restore_connector,)returninstance@classmethoddef_get_ngc_pretrained_model_info(cls,model_name:str,refresh_cache:bool=False)->Tuple[type,str]:""" Resolve the NGC model pretrained information given a model name. Assumes the model subclass implements the `list_available_models()` inherited method. Args: model_name: Str name of the model. Must be the original name or an alias of the model, without any '/'. refresh_cache: Bool, determines whether cache must be refreshed (model is re-downloaded). Returns: A tuple of details describing : - The resolved class of the model. This requires subclass to implement PretrainedModelInfo.class_. If the class cannot be resolved, default to the class that called this method. - The path to the NeMo model (.nemo file) in some cached directory. """location_in_the_cloud=Nonedescription=Noneclass_=Nonemodels=cls.list_available_models()ifmodelsisnotNone:forpretrained_model_infoincls.list_available_models():found=Falseifpretrained_model_info.pretrained_model_name==model_name:found=Trueelifpretrained_model_info.aliasesisnotNone:foraliasinpretrained_model_info.aliases:ifalias==model_name:found=Truebreakiffound:location_in_the_cloud=pretrained_model_info.locationdescription=pretrained_model_info.descriptionclass_=pretrained_model_info.class_breakiflocation_in_the_cloudisNone:raiseFileNotFoundError(f"Model {model_name} was not found. Check cls.list_available_models()\n"f"for the list of all available models.")# Use PurePosixPath for cloud URLs which always use forward slashesfilename=PurePosixPath(location_in_the_cloud).nameurl=location_in_the_cloud.replace(filename,"")cache_dir=Path.joinpath(resolve_cache_dir(),f'{filename[:-5]}')# If either description and location in the cloud changes, this will force re-downloadcache_subfolder=hashlib.md5((location_in_the_cloud+description).encode('utf-8')).hexdigest()# if file exists on cache_folder/subfolder, it will be re-used, unless refresh_cache is Truenemo_model_file_in_cache=maybe_download_from_cloud(url=url,filename=filename,cache_dir=cache_dir,subfolder=cache_subfolder,refresh_cache=refresh_cache)logging.info("Instantiating model from pre-trained checkpoint")ifclass_isNone:class_=clsreturnclass_,nemo_model_file_in_cache@classmethoddef_get_hf_hub_pretrained_model_info(cls,model_name:str,refresh_cache:bool=False)->Tuple[type,str]:""" Resolve the HuggingFace Hub model pretrained information given a model name. The model name must be of general syntax ``{source_repo}/{model_name}``. Note: The ``{source_repo}`` need not be ``nvidia``, it can be any public repository, even external to Nvidia. This allows public, externally contributed models to be run freely using Nvidia NeMo. Args: model_name: Str name of the model. Must be the original name or an alias of the model, without any '/'. refresh_cache: Bool, determines whether cache must be refreshed (model is re-downloaded). Returns: A tuple of details describing : - The resolved class of the model. Since the source is external to NeMo, always default to using the calling class. Depend on target class resolution by restore_from() for calling the correct class. - The path to the NeMo model (.nemo file) in some cached directory (managed by HF Hub). """# Resolve the model name without origin for filename# Use PurePosixPath since HuggingFace repo names use forward slashes (e.g., "nvidia/model-name")resolved_model_filename=PurePosixPath(model_name).name+'.nemo'# Try to take from cache first - if not fallback to options belowifnotrefresh_cache:path=try_to_load_from_cache(repo_id=model_name,filename=resolved_model_filename)ifpathisnotNoneandpathisnot_CACHED_NO_EXIST:returncls,path# Check if api token exists, use if it doeshf_token=get_hf_token()# First check if .nemo file exists in HFapi=HfApi(token=hf_token)# Check if model exists in HFnemo_file_exists=api.file_exists(repo_id=model_name,filename=resolved_model_filename,repo_type="model")ifnemo_file_exists:# Try to load the model from the Huggingface Hubpath=hf_hub_download(repo_id=model_name,filename=resolved_model_filename,library_name='nemo',library_version=nemo.__version__,force_download=refresh_cache,token=hf_token,)else:repo_info=api.repo_info(repo_id=model_name,token=hf_token,files_metadata=True)# Download whole HF repo and load entire directory as nemo directorycache_dir=Path.joinpath(resolve_cache_dir(),"hf_hub_cache",f'{model_name}')# If either description and location in the cloud changes, this will force re-downloadcache_subfolder=[]# Calculate hash of repo_infoforsiblinginrepo_info.siblings:filename=sibling.rfilename.lower()# Ignore updates to readme when downloading hashif"readme"notinfilenameor"git"notinfilename:cache_subfolder.append(sibling.blob_id)cache_subfolder=sorted(cache_subfolder)cache_subfolder="".join(cache_subfolder)cache_subfolder=hashlib.md5(cache_subfolder.encode('utf-8')).hexdigest()# if file exists on cache_folder/subfolder, it will be re-used, unless refresh_cache is Truesave_path=os.path.join(cache_dir,cache_subfolder)# If the cache dir already exists, delete it to preserve disk spaceifos.path.exists(cache_dir):num_files_in_dir=len(os.listdir(cache_dir))ifnum_files_in_dir>0:logging.info("Found {} files in cache directory {}".format(num_files_in_dir,cache_dir))logging.info(f"Deleting old cache directory for model `{model_name}` in order to prevent duplicates...")shutil.rmtree(cache_dir,ignore_errors=True)ifnotos.path.exists(save_path):logging.info(f"Downloading {model_name} from HuggingFace Hub to path: {save_path}")os.makedirs(save_path,exist_ok=True)path=snapshot_download(repo_id=model_name,library_name='nemo',library_version=nemo.__version__,force_download=refresh_cache,cache_dir=save_path,local_dir=save_path,local_dir_use_symlinks=False,token=hf_token,)returncls,pathdefgenerate_model_card(self,type:str="hf",template:str=None,template_kwargs:Optional[Dict[str,str]]=None)->object:""" Generates a ModelCard for the current model. This method is called when pushing the model to the Hub. Returns: An object that can be represented as a str representation of the model card, usually in Markdown format. """iftemplateisNone:template=copy.deepcopy(NEMO_DEFAULT_MODEL_CARD_TEMPLATE)# Populate template kwargs with common model card fieldsiftemplate_kwargsisNone:template_kwargs={}iftype=="hf":# Use HuggingFaceFileIO method to generate the huggingface model cardreturnself._get_hf_model_card(template=template,template_kwargs=template_kwargs)else:raiseValueError(f"Model card type {type} not supported.")
classtypecheck:""" A decorator which performs input-output neural type checks, and attaches neural types to the output of the function that it wraps. Requires that the class inherit from :class:`~nemo.core.Typing` in order to perform type checking, and will raise an error if that is not the case. # Usage (Class level type support) .. code-block:: python @typecheck() def fn(self, arg1, arg2, ...): ... # Usage (Function level type support) .. code-block:: python @typecheck(input_types=..., output_types=...) def fn(self, arg1, arg2, ...): ... Points to be noted: 1) The brackets () in `@typecheck()` are necessary. You will encounter a TypeError: __init__() takes 1 positional argument but X were given without those brackets. 2) The function can take any number of positional arguments during definition. When you call this function, all arguments must be passed using kwargs only. """
classTypeState(Enum):""" Placeholder to denote the default value of type information provided. If the constructor of this decorator is used to override the class level type definition, this enum value indicate that types will be overridden. """UNINITIALIZED=0
def__init__(self,input_types:Union[TypeState,Dict[str,NeuralType]]=TypeState.UNINITIALIZED,output_types:Union[TypeState,Dict[str,NeuralType]]=TypeState.UNINITIALIZED,ignore_collections:bool=False,):self.input_types=input_typesself.output_types=output_typesifinput_types==self.TypeState.UNINITIALIZED:self.input_override=Falseelse:self.input_override=Trueifoutput_types==self.TypeState.UNINITIALIZED:self.output_override=Falseelse:self.output_override=Trueself.ignore_collections=ignore_collections
def__call__(self,wrapped):returnself.wrapped_call(wrapped)
defunwrapped_call(self,wrapped):"""Call without typechecking"""returnwrapped
@wrapt.decorator(enabled=is_typecheck_enabled)defwrapped_call(self,wrapped,instance:Typing,args,kwargs):""" Wrapper method that can be used on any function of a class that implements :class:`~nemo.core.Typing`. By default, it will utilize the `input_types` and `output_types` properties of the class inheriting Typing. Local function level overrides can be provided by supplying dictionaries as arguments to the decorator. Args: input_types: Union[TypeState, Dict[str, NeuralType]]. By default, uses the global `input_types`. output_types: Union[TypeState, Dict[str, NeuralType]]. By default, uses the global `output_types`. ignore_collections: Bool. Determines if container types should be asserted for depth checks, or if depth checks are skipped entirely. """ifinstanceisNone:raiseRuntimeError("Only classes which inherit nemo.core.Typing can use this decorator !")ifnotisinstance(instance,Typing):raiseRuntimeError("Only classes which inherit nemo.core.Typing can use this decorator !")ifhasattr(instance,'input_ports')orhasattr(instance,'output_ports'):raiseRuntimeError("Typing requires override of `input_types()` and `output_types()`, ""not `input_ports() and `output_ports()`")# Preserve type informationifself.input_typesistypecheck.TypeState.UNINITIALIZED:self.input_types=instance.input_typesifself.output_typesistypecheck.TypeState.UNINITIALIZED:self.output_types=instance.output_types# Resolve global type or local overridden typeifself.input_override:input_types=self.input_typeselse:input_types=instance.input_typesifself.output_override:output_types=self.output_typeselse:output_types=instance.output_types# If types are not defined, skip type checks and just call the wrapped methodifinput_typesisNoneandoutput_typesisNone:returnwrapped(*args,**kwargs)# Check that all arguments are kwargsifinput_typesisnotNoneandlen(args)>0:raiseTypeError("All arguments must be passed by kwargs only for typed methods")# Perform rudimentary input checks hereinstance._validate_input_types(input_types=input_types,ignore_collections=self.ignore_collections,**kwargs)# Call the method - this can be forward, or any other callable methodoutputs=wrapped(*args,**kwargs)instance._attach_and_validate_output_types(output_types=output_types,ignore_collections=self.ignore_collections,out_objects=outputs)returnoutputs
@staticmethoddefset_typecheck_enabled(enabled:bool=True):""" Global method to enable/disable typechecking. Args: enabled: bool, when True will enable typechecking. """global_TYPECHECK_ENABLED_TYPECHECK_ENABLED=enabled
@staticmethod@contextmanagerdefdisable_checks():""" Context manager that temporarily disables type checking within its context. """typecheck.set_typecheck_enabled(enabled=False)try:yieldfinally:typecheck.set_typecheck_enabled(enabled=True)
@staticmethoddefset_semantic_check_enabled(enabled:bool=True):""" Global method to enable/disable semantic typechecking. Args: enabled: bool, when True will enable semantic typechecking. """global_TYPECHECK_SEMANTIC_CHECK_ENABLED_TYPECHECK_SEMANTIC_CHECK_ENABLED=enabled
@staticmethod@contextmanagerdefdisable_semantic_checks():""" Context manager that temporarily disables semantic type checking within its context. """typecheck.set_semantic_check_enabled(enabled=False)try:yieldfinally:typecheck.set_semantic_check_enabled(enabled=True)
@staticmethoddefenable_wrapping(enabled:bool=True):"""Enables typechecking"""typecheck.set_typecheck_enabled(enabled)ifenabled:typecheck.__call__=nemo.core.classes.common.typecheck.wrapped_callelse:typecheck.__call__=nemo.core.classes.common.typecheck.unwrapped_call