import warnings from time import sleep from typing import Any, Dict, List, Optional, Tuple, Union import torch from omegaconf import DictConfig from transformers.cache_utils import DynamicCache from src.utils import pylogger, rich_utils from functools import partial from tokenizers.processors import TemplateProcessing from src.data.loading.components.interfaces import TokenizerConfig log = pylogger.RankedLogger(__name__, rank_zero_only=True) def print_warnings_for_missing_configs(cfg: DictConfig) -> None: _DEFAULT_CONFIGS = [ "data_loading", "model", "loss", "optim", "eval", ] has_warnings = False for config in _DEFAULT_CONFIGS: if not cfg.get(config): log.warning( f"Config {config} was not found in the config tree. Make sure this is expected." ) has_warnings = True if has_warnings: sleep(3) # wait for 3 seconds to let the user read the warning def extras(cfg: DictConfig) -> None: """Applies optional utilities before the task is started. Utilities: - Ignoring python warnings - Setting tags from command line - Rich config printing :param cfg: A DictConfig object containing the config tree. """ # return if no `extras` config if not cfg.get("extras"): log.warning("Extras config not found! ") return # disable python warnings if cfg.extras.get("ignore_warnings"): log.info("Disabling python warnings! ") warnings.filterwarnings("ignore") # prompt user to input tags from command line if none are provided in the config if cfg.extras.get("enforce_tags"): log.info("Enforcing tags! ") rich_utils.enforce_tags(cfg, save_to_file=True) if cfg.extras.get("print_config_warnings"): print_warnings_for_missing_configs(cfg) # pretty print config tree using Rich library if cfg.extras.get("print_config"): log.info("Printing config tree with Rich! ") rich_utils.print_config_tree(cfg, resolve=True, save_to_file=True) def delete_module(module: torch.nn.Module, module_name: str) -> None: """Recursively delete a submodule from a module. :param module: the parent module that we want the submodule to be removed from. :param module_name: the name of the submodule to be removed. :return: None. """ if hasattr(module, module_name): delattr(module, module_name) for name, submodule in module.named_children(): delete_module(submodule, module_name) def find_module_shape( module: torch.nn.Module, module_name: str ) -> Optional[torch.Size]: """Recursively find a submodule in a module and return its shape. :param module: the parent module that we want the submodule to be removed from. :param module_name: the name of the submodule to be removed. :return: the shape of the module if it exists. """ if hasattr(module, module_name): return getattr(module, module_name).weight.shape for name, submodule in module.named_children(): shape = find_module_shape(submodule, module_name) if shape: return shape return None def reset_parameters(module: torch.nn.Module) -> None: """Reset the parameters of a given module. :param module: the module whose parameters will be reset. :return: None. """ if hasattr(module, "reset_parameters"): module.reset_parameters() else: for layer in module.children(): reset_parameters(layer) def get_var_if_not_none(value: Optional[Any], default_value: Any) -> Any: """ :return: value if value is not None, else default_value Note that when value is: Boolean: False is not None Int: 0 is not None List: An empty list is not None Tensor: A tensor with all zeros is not None """ return value if value is not None else default_value def get_class_name_str(class_definition: Any) -> str: """ Args: class_definition: The class definition. Returns: The fully qualified name of the given class. """ return ".".join([class_definition.__module__, class_definition.__name__]) def has_class_object_inside_list(obj_list: list, class_type: Any) -> bool: """ Args: obj_list: List of objects. class_type: The class type to check. Returns: True if the list contains an object of the given class type. """ return any(isinstance(obj, class_type) for obj in obj_list) def convert_legacy_kv_cache_to_dynamic( kv_cache: Union[DynamicCache, Tuple[torch.Tensor]] ) -> DynamicCache: """ Converts a legacy key-value cache (Tuple of tensors) to a dynamic cache. Args: kv_cache: The key-value cache. Returns: The dynamic cache. """ if isinstance(kv_cache, DynamicCache): return kv_cache return DynamicCache.from_legacy_cache(kv_cache) def get_parent_module_and_attr( model: torch.nn.Module, module_name: str ) -> Tuple[torch.nn.Module, str]: """ Get the parent module and attribute name for a given module name. Args: model (torch.nn.Module): The model containing the module. module_name (str): The full name of the module. Returns: Tuple[torch.nn.Module, str]: The parent module and the attribute name. """ parts = module_name.split(".") parent = model for part in parts[:-1]: parent = getattr(parent, part) return parent, parts[-1] def lightning_precision_to_dtype(precision: str) -> torch.dtype: """ Convert Lightning precision identifier to PyTorch dtype. Args: precision (str): The precision identifier used in Lightning. Expected values include '32', '32-true', '16', '16-mixed', 'bf16', '64', 'half'. Returns: torch.dtype: The corresponding PyTorch dtype. Raises: ValueError: If an unsupported precision type is provided. """ # Mapping from Lightning precision identifiers to PyTorch dtypes precision_map = { "32": torch.float32, # Single precision (float32) "32-true": torch.float32, # Also maps to float32, useful for clarity when specifying defaults "64": torch.float64, # Double precision "16": torch.float16, # Half precision "16-mixed": torch.float16, # Mixed precision typically uses torch.float16 "bf16": torch.bfloat16, # BFloat16 precision "half": torch.float16, # Alias for half precision } if precision in precision_map: return precision_map[precision] else: raise ValueError( f"Unsupported precision type: '{precision}'. " "Supported precision types are: '32', '32-true', '64', '16', '16-mixed', 'bf16', 'half'." ) def load_tokenize(config: TokenizerConfig) -> Any: """Load tokenizer and return a partial function for tokenization.""" tokenizer = config.tokenizer if hasattr(config, "special_tokens"): tokenizer.add_special_tokens(config.special_tokens) if config.postprocess_eos_token: tokenizer._tokenizer.post_processor = TemplateProcessing( single="$A " + tokenizer.eos_token, special_tokens=[(tokenizer.eos_token, tokenizer.eos_token_id)], ) tokenize = partial( tokenizer.encode_plus, max_length=config.max_length, padding=config.padding, truncation=config.truncation, add_special_tokens=config.add_special_tokens, return_tensors="pt", ) return tokenize def sample_gumbel(shape: Tuple, device: torch.device, eps=1e-20) -> torch.Tensor: """Sample from Gumbel(0, 1)""" U = torch.rand(shape, device=device) return -torch.log(-torch.log(U + eps) + eps) def gumbel_softmax_sample(logits: torch.Tensor, temperature: float) -> torch.Tensor: """Draw a sample from the Gumbel-Softmax distribution""" y = logits + sample_gumbel(logits.shape, logits.device) sample = F.softmax(y / temperature, dim=-1) return sample