import contextlib import logging from abc import ABCMeta, abstractmethod from collections import OrderedDict from typing import Any import numpy as np import torch import torchmetrics from ludwig.combiners.combiners import Combiner from ludwig.constants import COMBINED, LOSS, NAME from ludwig.encoders.base import Encoder from ludwig.features.base_feature import InputFeature, NonPropertyModuleWrapper, OutputFeature from ludwig.features.feature_registries import get_input_type_registry, get_output_type_registry from ludwig.features.feature_utils import LudwigFeatureDict from ludwig.features.passthrough_feature import create_passthrough_input_feature from ludwig.modules.metric_modules import LudwigMetric from ludwig.modules.training_hooks import TrainingHook from ludwig.schema.features.base import BaseInputFeatureConfig, BaseOutputFeatureConfig, FeatureCollection from ludwig.utils.algorithms_utils import topological_sort_feature_dependencies from ludwig.utils.metric_utils import get_scalar_from_ludwig_metric from ludwig.utils.misc_utils import get_from_registry from ludwig.utils.torch_utils import LudwigModule, reg_loss logger = logging.getLogger(__name__) class BaseModel(LudwigModule, metaclass=ABCMeta): """Base model for use in LudwigModule. Implementations of this class should implement the following methods: - type() - forward() """ @staticmethod @abstractmethod def type() -> str: """Returns the model type.""" def __init__(self, random_seed: int | None = None): self._random_seed = random_seed # Ensures model weight initialization is deterministic even though set_random_seed() # is called later in the trainer. Required for test_api::test_api_training_determinism. if random_seed is not None: torch.random.manual_seed(random_seed) super().__init__() self.input_features = self.create_feature_dict() self.output_features = self.create_feature_dict() # ================ Combined loss metric ================ self._eval_loss_metric = NonPropertyModuleWrapper(torchmetrics.MeanMetric()) self._eval_additional_losses_metrics = NonPropertyModuleWrapper(torchmetrics.MeanMetric()) # ================ Training Hook Handles ================ self._forward_hook_handles: list[TrainingHook] = [] def create_feature_dict(self) -> LudwigFeatureDict: """Creates and returns a LudwigFeatureDict.""" return LudwigFeatureDict() def to_device(self, device): return self.to(device) def metrics_to_device(self, device: str): self._eval_loss_metric.module = self._eval_loss_metric.module.to(device) self._eval_additional_losses_metrics.module = self._eval_additional_losses_metrics.module.to(device) for feature in self.output_features.values(): feature._eval_loss_metric.module = feature._eval_loss_metric.module.to(device) @classmethod def build_inputs(cls, input_feature_configs: FeatureCollection[BaseInputFeatureConfig]) -> dict[str, InputFeature]: """Builds and returns input features in topological order.""" input_features = OrderedDict() input_features_def = topological_sort_feature_dependencies(input_feature_configs.to_list()) for input_feature_def in input_features_def: input_features[input_feature_def[NAME]] = cls.build_single_input( getattr(input_feature_configs, input_feature_def[NAME]), input_features ) return input_features @staticmethod def build_single_input( feature_config: BaseInputFeatureConfig, other_input_features: dict[str, InputFeature] | None ) -> InputFeature: """Builds a single input feature from the input feature definition.""" logger.debug(f"Input {feature_config.type} feature {feature_config.name}") encoder_obj = None if feature_config.tied is not None: tied_input_feature_name = feature_config.tied if tied_input_feature_name in other_input_features: encoder_obj = other_input_features[tied_input_feature_name].encoder_obj return create_input_feature(feature_config, encoder_obj) @classmethod def build_outputs( cls, output_feature_configs: FeatureCollection[BaseOutputFeatureConfig], combiner: Combiner ) -> dict[str, OutputFeature]: """Builds and returns output features in topological order.""" output_features_def = topological_sort_feature_dependencies(output_feature_configs.to_list()) output_features = {} for output_feature_def in output_features_def: # TODO(Justin): Check that the semantics of input_size align with what the combiner's output shape returns # for seq2seq. getattr(output_feature_configs, output_feature_def[NAME]).input_size = combiner.output_shape[-1] output_features[output_feature_def[NAME]] = cls.build_single_output( getattr(output_feature_configs, output_feature_def[NAME]), output_features ) return output_features @staticmethod def build_single_output( feature_config: BaseOutputFeatureConfig, output_features: dict[str, OutputFeature] | None ) -> OutputFeature: """Builds a single output feature from the output feature definition.""" logger.debug(f"Output {feature_config.type} feature {feature_config.name}") output_feature_class = get_from_registry(feature_config.type, get_output_type_registry()) output_feature_obj = output_feature_class(feature_config, output_features=output_features) return output_feature_obj def get_model_inputs(self): """Returns a dict of feature name -> sample model input.""" device = next(self.parameters()).device inputs = { input_feature_name: input_feature.create_sample_input().to(device) for input_feature_name, input_feature in self.input_features.items() } return inputs def get_model_size(self) -> int: """Returns total number of parameters in model.""" model_tensors = self.collect_weights() total_size = 0 for tnsr in model_tensors: total_size += tnsr[1].detach().cpu().numpy().size return total_size @property def input_shape(self) -> torch.Size: """Returns the shape of a single model input (excluding batch dimension). Subclasses should override this to return a meaningful shape. The default is a (1,) scalar — sufficient for models that don't rely on input_shape for decoder sizing. """ return torch.Size([1]) @abstractmethod def forward( self, inputs: ( dict[str, torch.Tensor] | dict[str, np.ndarray] | tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]] ), mask=None, ) -> dict[str, torch.Tensor]: """Forward pass of the model. Args: inputs: Inputs to the model. Can be a dictionary of input names to input tensors or a tuple of (inputs, targets) where inputs is a dictionary of input names to input tensors and targets is a dictionary of target names to target tensors. mask: A mask for the inputs. Returns: A dictionary of output {feature name}::{tensor_name} -> output tensor. """ def predictions(self, inputs): """Returns the model's predictions for the given inputs.""" outputs = self(inputs) return self.outputs_to_predictions(outputs) def outputs_to_predictions(self, outputs: dict[str, torch.Tensor]) -> dict[str, dict[str, torch.Tensor]]: """Returns the model's predictions given the raw model outputs.""" predictions = {} for of_name in self.output_features: predictions[of_name] = self.output_features.get(of_name).predictions(outputs, of_name) return predictions def evaluation_step(self, inputs, targets): """Predict the inputs and update evaluation metrics.""" predictions = self.predictions(inputs) self.update_metrics(targets, predictions) return predictions def predict_step(self, inputs): """Predict the inputs.""" return self.predictions(inputs) def train_loss( self, targets, predictions, regularization_type: str | None = None, regularization_lambda: float | None = None, ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: """Computes the training loss for the model. Args: targets: A dictionary of target names to target tensors. predictions: A dictionary of output names to output tensors. regularization_type: One of 'l1', 'l2', 'l1_l2', or None. regularization_lambda: The regularization lambda. Returns: A tuple of the loss tensor and a dictionary of loss for every output feature. """ of_train_losses = {} of_weights = {} for of_name, of_obj in self.output_features.items(): of_train_loss = of_obj.train_loss(targets[of_name], predictions, of_name) of_train_losses[of_name] = of_train_loss of_weights[of_name] = of_obj.loss.weight # Use loss balancer if available, otherwise static weighted sum if hasattr(self, "loss_balancer") and self.loss_balancer is not None: train_loss = self.loss_balancer(of_train_losses, of_weights) else: train_loss = sum(of_weights[k] * of_train_losses[k] for k in of_train_losses) additional_losses = self.losses() if additional_losses: train_loss += torch.sum(torch.stack(additional_losses)) # other losses # Add regularization loss if regularization_type is not None and regularization_lambda != 0: train_loss += reg_loss(self, regularization_type, l1=regularization_lambda, l2=regularization_lambda) return train_loss, of_train_losses def eval_loss(self, targets, predictions): """Computes all evaluation losses for the model given targets and predictions. Args: targets: A dictionary of target names to target tensors. predictions: A dictionary of output names to output tensors. Returns: A tuple of loss values for eval losses and additional losses. """ eval_loss = 0 for of_name, of_obj in self.output_features.items(): of_eval_loss = of_obj.eval_loss(targets[of_name], predictions[of_name]) eval_loss += of_obj.loss.weight * of_eval_loss additional_loss = 0 additional_losses = self.losses() if additional_losses: additional_loss = torch.sum(torch.stack(additional_losses)) # other losses return eval_loss, additional_loss def update_metrics(self, targets, predictions): """Updates the model's metrics given targets and predictions.""" for of_name, of_obj in self.output_features.items(): of_obj.update_metrics(targets[of_name], predictions[of_name]) eval_loss, additional_losses = self.eval_loss(targets, predictions) self.eval_loss_metric.update(eval_loss) self.eval_additional_losses_metrics.update(additional_losses) @property def eval_loss_metric(self) -> LudwigMetric: return self._eval_loss_metric.module @eval_loss_metric.setter def eval_loss_metric(self, value: LudwigMetric) -> None: self._eval_loss_metric.module = value @property def eval_additional_losses_metrics(self) -> LudwigMetric: return self._eval_additional_losses_metrics.module def get_metrics(self) -> dict[str, dict[str, float]]: """Returns a dictionary of metrics for each output feature of the model.""" all_of_metrics = {} for of_name, of_obj in self.output_features.items(): all_of_metrics[of_name] = of_obj.get_metrics() all_of_metrics[COMBINED] = { LOSS: get_scalar_from_ludwig_metric(self.eval_loss_metric) + get_scalar_from_ludwig_metric(self.eval_additional_losses_metrics) } return all_of_metrics def reset_metrics(self): """Resets the model's metrics.""" for of_obj in self.output_features.values(): of_obj.reset_metrics() self.eval_loss_metric.reset() def collect_weights(self, tensor_names=None, **kwargs): """Returns named parameters filtered against `tensor_names` if not None.""" if not tensor_names: return self.named_parameters() # Check for bad tensor names. weight_names = {name for name, _ in self.named_parameters()} for name in tensor_names: if name not in weight_names: raise ValueError(f'Requested tensor name filter "{name}" not present in the model graph') # Apply filter. tensor_set = set(tensor_names) return [named_param for named_param in self.named_parameters() if named_param[0] in tensor_set] def unskip(self): """Converts all skipped features into their fully encoded versions.""" @abstractmethod def save(self, save_path: str): """Saves the model to the given path.""" @abstractmethod def load(self, save_path: str): """Loads the model from the given path.""" @abstractmethod def get_args(self): """Returns init arguments for constructing this model.""" @contextlib.contextmanager def use_generation_config(self, generation_config: dict[str, Any]): if generation_config is not None: raise NotImplementedError(f"{self.__class__.__name__} does not support generation_config. ") yield def _activate_forward_hooks(self): """Activates/registers forward hooks for the model.""" def _deactivate_forward_hooks(self) -> None: """Deactivates/de-registers forward hooks for the model (if needed).""" for handle in self._forward_hook_handles: handle.deactivate_hook() def create_input_feature(feature_config: BaseInputFeatureConfig, encoder_obj: Encoder | None) -> InputFeature: input_feature_cls = get_from_registry(feature_config.type, get_input_type_registry()) input_feature = input_feature_cls(feature_config, encoder_obj=encoder_obj) if not feature_config.encoder.skip: return input_feature return create_passthrough_input_feature(input_feature, feature_config)