tzrec/models/hstu.py

# Copyright (c) 2024-2025, Alibaba Group;
# 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.

from typing import Any, Dict, List, Optional

import torch
import torch.nn.functional as F
from torch._tensor import Tensor

from tzrec.datasets.utils import NEG_DATA_GROUP, Batch
from tzrec.features.feature import BaseFeature
from tzrec.models.match_model import MatchModel, MatchTowerWoEG
from tzrec.modules.embedding import EmbeddingGroup
from tzrec.modules.sequence import HSTUEncoder
from tzrec.protos import model_pb2, simi_pb2, tower_pb2
from tzrec.utils import config_util


@torch.fx.wrap
def _update_dict_tensor(
    tensor_dict: Dict[str, torch.Tensor],
    new_tensor_dict: Optional[Dict[str, Optional[torch.Tensor]]],
) -> None:
    if new_tensor_dict:
        for k, v in new_tensor_dict.items():
            if v is not None:
                tensor_dict[k] = v


class HSTUMatchUserTower(MatchTowerWoEG):
    """HSTU Match model user tower.

    Args:
        tower_config (Tower): user tower config.
        output_dim (int): user output embedding dimension.
        similarity (Similarity): when use COSINE similarity,
            will norm the output embedding.
        feature_group (FeatureGroupConfig): feature group config.
        features (list): list of features.
    """

    def __init__(
        self,
        tower_config: tower_pb2.HSTUMatchTower,
        output_dim: int,
        similarity: simi_pb2.Similarity,
        feature_group: model_pb2.FeatureGroupConfig,
        feature_group_dims: List[int],
        features: List[BaseFeature],
        model_config: model_pb2.ModelConfig,
    ) -> None:
        super().__init__(tower_config, output_dim, similarity, feature_group, features)
        self.tower_config = tower_config
        encoder_config = tower_config.hstu_encoder
        seq_config_dict = config_util.config_to_kwargs(encoder_config)
        sequence_dim = sum(feature_group_dims)
        seq_config_dict["sequence_dim"] = sequence_dim
        self.seq_encoder = HSTUEncoder(**seq_config_dict)

    def forward(self, grouped_features: Dict[str, torch.Tensor]) -> torch.Tensor:
        """Forward the tower.

        Args:
            grouped_features: Dictionary containing grouped feature tensors

        Returns:
            torch.Tensor: The output tensor from the tower
        """
        output = self.seq_encoder(grouped_features)

        return output


class HSTUMatchItemTower(MatchTowerWoEG):
    """HSTU Match model item tower.

    Args:
        tower_config (Tower): item tower config.
        output_dim (int): item output embedding dimension.
        similarity (Similarity): when use COSINE similarity,
            will norm the output embedding.
        feature_group (FeatureGroupConfig): feature group config.
        features (list): list of features.
    """

    def __init__(
        self,
        tower_config: tower_pb2.Tower,
        output_dim: int,
        similarity: simi_pb2.Similarity,
        feature_group: model_pb2.FeatureGroupConfig,
        features: List[BaseFeature],
    ) -> None:
        super().__init__(tower_config, output_dim, similarity, feature_group, features)
        self.tower_config = tower_config

    def forward(self, grouped_features: Dict[str, torch.Tensor]) -> torch.Tensor:
        """Forward the tower.

        Args:
            grouped_features: Dictionary containing grouped feature tensors

        Returns:
            torch.Tensor: The output tensor from the tower
        """
        output = grouped_features[f"{self._group_name}.sequence"]
        output = F.normalize(output, p=2.0, dim=1, eps=1e-6)

        return output


class HSTUMatch(MatchModel):
    """HSTU Match model.

    Args:
        model_config (ModelConfig): an instance of ModelConfig.
        features (list): list of features.
        labels (list): list of label names.
    """

    def __init__(
        self,
        model_config: model_pb2.ModelConfig,
        features: List[BaseFeature],
        labels: List[str],
        sample_weights: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        super().__init__(model_config, features, labels, sample_weights, **kwargs)
        assert len(model_config.feature_groups) == 1
        self.embedding_group = EmbeddingGroup(
            features, list(model_config.feature_groups)
        )

        name_to_feature_group = {x.group_name: x for x in model_config.feature_groups}
        feature_group = name_to_feature_group[self._model_config.hstu_tower.input]

        used_features = self.get_features_in_feature_groups([feature_group])

        self.user_tower = HSTUMatchUserTower(
            self._model_config.hstu_tower,
            self._model_config.output_dim,
            self._model_config.similarity,
            feature_group,
            self.embedding_group.group_dims(
                self._model_config.hstu_tower.input + ".sequence"
            ),
            used_features,
            model_config,
        )

        self.item_tower = HSTUMatchItemTower(
            self._model_config.hstu_tower,
            self._model_config.output_dim,
            self._model_config.similarity,
            feature_group,
            used_features,
        )

        self.seq_tower_input = self._model_config.hstu_tower.input

    def predict(self, batch: Batch) -> Dict[str, Tensor]:
        """Forward the model.

        Args:
            batch (Batch): input batch data.

        Return:
            predictions (dict): a dict of predicted result.
        """
        batch_sparse_features = batch.sparse_features[NEG_DATA_GROUP]
        # Get batch_size and neg_sample_size from batch_sparse_features
        batch_size = batch.labels[self._label_name].shape[0]
        neg_sample_size = batch_sparse_features.lengths()[batch_size] - 1
        grouped_features = self.embedding_group(batch)

        item_group_features = {
            self.seq_tower_input + ".sequence": grouped_features[
                self.seq_tower_input + ".sequence"
            ][batch_size:, : neg_sample_size + 1],
        }
        item_tower_emb = self.item_tower(item_group_features)
        user_group_features = {
            self.seq_tower_input + ".sequence": grouped_features[
                self.seq_tower_input + ".sequence"
            ][:batch_size],
            self.seq_tower_input + ".sequence_length": grouped_features[
                self.seq_tower_input + ".sequence_length"
            ][:batch_size],
        }
        user_tower_emb = self.user_tower(user_group_features)
        ui_sim = (
            self.sim(user_tower_emb, item_tower_emb, neg_for_each_sample=True)
            / self._model_config.temperature
        )
        return {"similarity": ui_sim}

    def sim(
        self,
        user_emb: torch.Tensor,
        item_emb: torch.Tensor,
        neg_for_each_sample: bool = False,
    ) -> torch.Tensor:
        """Override the sim method in MatchModel to calculate similarity."""
        if self._in_batch_negative:
            return torch.mm(user_emb, item_emb.T)
        else:
            batch_size = user_emb.size(0)
            pos_item_emb = item_emb[:, 0]
            neg_item_emb = item_emb[:, 1:].reshape(-1, item_emb.shape[-1])
            pos_ui_sim = torch.sum(
                torch.multiply(user_emb, pos_item_emb), dim=-1, keepdim=True
            )
            neg_ui_sim = None
            if not neg_for_each_sample:
                neg_ui_sim = torch.matmul(user_emb, neg_item_emb.transpose(0, 1))
            else:
                num_neg_per_user = neg_item_emb.size(0) // batch_size
                neg_size = batch_size * num_neg_per_user
                neg_item_emb = neg_item_emb[:neg_size]
                neg_item_emb = neg_item_emb.view(batch_size, num_neg_per_user, -1)
                neg_ui_sim = torch.sum(user_emb.unsqueeze(1) * neg_item_emb, dim=-1)
            return torch.cat([pos_ui_sim, neg_ui_sim], dim=-1)