[docs]defsupported_tasks()->List[str]:"Returns the list of supported tasks."returnsorted(mod.__TASK__formodin__TASKS__)
[docs]defreduce_model_config(config:Any,task:str)->Dict[str,Any]:"""Reduces a model size."""head_size0=(config.head_dimifhasattr(config,"head_dim")andconfig.head_dimelse(config.hidden_size//config.num_attention_headsifhasattr(config,"hidden_size")andhasattr(config,"num_attention_heads")elseNone))tasks={mod.__TASK__:mod.reduce_model_configformodin__TASKS__}asserttaskintasks,f"Task {task!r} not found in {sorted(tasks)}"res=tasks[task](config)ifhead_size0and"head_dim"inres:head_size=(config.head_dimifhasattr(config,"head_dim")andconfig.head_dimelseconfig.hidden_size//config.num_attention_heads)asserthead_size0==head_sizeorhead_size%16==0,(f"head_size should be a multiple of 16 "f"(head_size0={head_size0}), res={res}, "f"config=\n{config}")returnres
[docs]defrandom_input_kwargs(config:Any,task:str)->Tuple[Dict[str,Any],Callable]:""" Inputs kwargs. If the configuration is None, the function selects typical dimensions. It returns parameters and a function. The function creates dummy inputs if it receives the parameters returned as a first result. """tasks={mod.__TASK__:mod.random_input_kwargsformodin__TASKS__}asserttaskintasks,f"Task {task!r} not found in {sorted(tasks)}"returntasks[task](config)