/ludwig/models/base.py

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)