From 174df0f7ae0f8552876cc2289630b8adfb6f1d71 Mon Sep 17 00:00:00 2001 From: yangzhou23 Date: Tue, 21 Apr 2026 18:00:35 +0800 Subject: [PATCH] feat(model): add gru4rec --- README.md | 4 +- README_en.md | 4 +- nextrec/__version__.py | 2 +- nextrec/models/ranking/bst.py | 16 +-- nextrec/models/ranking/dlrm.py | 2 +- nextrec/models/sequential/base.py | 4 +- nextrec/models/sequential/gru4rec.py | 162 +++++++++++++++++++++++++ pyproject.toml | 2 +- test/test_generative_models.py | 56 +++++++++ tutorials/run_all_sequential_models.py | 13 +- 10 files changed, 244 insertions(+), 21 deletions(-) create mode 100644 nextrec/models/sequential/gru4rec.py diff --git a/README.md b/README.md index 19e66a1..5bfbd10 100644 --- a/README.md +++ b/README.md @@ -9,7 +9,7 @@ ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg) ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/zerolovesea/NextRec/blob/main/tutorials/notebooks/zh/%E5%BF%AB%E9%80%9F%E5%85%A5%E9%97%A8nextrec.ipynb) -![Version](https://img.shields.io/badge/Version-0.6.12-orange.svg) +![Version](https://img.shields.io/badge/Version-0.6.13-orange.svg) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec) @@ -181,7 +181,7 @@ nextrec --mode=predict --predict_config=path/to/predict_config.yaml ## 兼容平台 -当前最新版本为0.6.12，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本： +当前最新版本为0.6.13，所有模型和测试代码均已在以下平台通过验证，如果开发者在使用中遇到兼容问题，请在issue区提出错误报告及系统版本： | 平台 | 配置 | |------|------| diff --git a/README_en.md b/README_en.md index 509a240..fcbaca6 100644 --- a/README_en.md +++ b/README_en.md @@ -8,7 +8,7 @@ ![Python](https://img.shields.io/badge/Python-3.10+-blue.svg) ![PyTorch](https://img.shields.io/badge/PyTorch-1.10+-ee4c2c.svg) ![License](https://img.shields.io/badge/License-Apache%202.0-green.svg) -![Version](https://img.shields.io/badge/Version-0.6.12-orange.svg) +![Version](https://img.shields.io/badge/Version-0.6.13-orange.svg) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/zerolovesea/NextRec/blob/main/tutorials/notebooks/en/Hands%20on%20nextrec.ipynb) [![Ask DeepWiki](https://deepwiki.com/badge.svg)](https://deepwiki.com/zerolovesea/NextRec) @@ -190,7 +190,7 @@ Prediction outputs are saved under `{checkpoint_path}/predictions/{name}.{save_d ## Platform Compatibility -The current version is 0.6.12. All models and test code have been validated on the following platforms. If you encounter compatibility issues, please report them in the issue tracker with your system version: +The current version is 0.6.13. All models and test code have been validated on the following platforms. If you encounter compatibility issues, please report them in the issue tracker with your system version: | Platform | Configuration | |----------|---------------| diff --git a/nextrec/__version__.py b/nextrec/__version__.py index 49aadb8..c98d6e3 100644 --- a/nextrec/__version__.py +++ b/nextrec/__version__.py @@ -1 +1 @@ -__version__ = "0.6.12" +__version__ = "0.6.13" diff --git a/nextrec/models/ranking/bst.py b/nextrec/models/ranking/bst.py index 9dcd81d..b94bd01 100644 --- a/nextrec/models/ranking/bst.py +++ b/nextrec/models/ranking/bst.py @@ -148,20 +148,14 @@ def __init__( int(feature.embedding_dim) != self.embed_dim for feature in (*self.behavior_features, *self.candidate_features) ): - raise ValueError( - "[BST Error] behavior and candidate features must share the same embedding_dim." - ) + raise ValueError("[BST Error] behavior and candidate features must share the same embedding_dim.") if self.embed_dim % num_heads != 0: raise ValueError( f"[BST Error] embedding_dim({self.embed_dim}) must be divisible by num_heads({num_heads})." ) exclude_feature_names = {feature.name for feature in (*self.behavior_features, *self.candidate_features)} - self.other_features = [ - feature - for feature in self.all_features - if feature.name not in exclude_feature_names - ] + self.other_features = [feature for feature in self.all_features if feature.name not in exclude_feature_names] self.embedding = EmbeddingLayer(features=self.all_features) encoder_layer = nn.TransformerEncoderLayer( @@ -175,7 +169,7 @@ def __init__( other_dim = self.embedding.compute_output_dim(self.other_features) target_dim = self.embedding.compute_output_dim(self.candidate_features) transformer_dim = len(self.behavior_features) * self.embed_dim - + # behavior fusion [B, Behavior_num * D] + target flatten [B, Target_num * D] + optional context mlp_input_dim = transformer_dim + target_dim + other_dim mlp_params.setdefault("hidden_dims", [256, 128]) @@ -189,7 +183,7 @@ def __init__( def forward(self, x) -> torch.Tensor: # candidate sparse features -> [B, Target_num, D] embed_x_target = self.embedding(x=x, features=self.candidate_features, squeeze_dim=False) - + # concat behavior sequence and target embeddings transformer_pooling = [] for behavior_feature in self.behavior_features: @@ -218,5 +212,5 @@ def forward(self, x) -> torch.Tensor: # auxiliary context features -> [B, Context_dim] features.append(self.embedding(x=x, features=self.other_features, squeeze_dim=True)) logits = self.mlp(torch.cat(features, dim=1)) - + return logits diff --git a/nextrec/models/ranking/dlrm.py b/nextrec/models/ranking/dlrm.py index e85e55a..a77e02c 100644 --- a/nextrec/models/ranking/dlrm.py +++ b/nextrec/models/ranking/dlrm.py @@ -192,5 +192,5 @@ def forward(self, x) -> torch.Tensor: # standard DLRM keeps the processed dense embedding alongside interactions top_inputs.insert(0, dense_embedding) logits = self.top_mlp(torch.cat(top_inputs, dim=1)) - + return logits diff --git a/nextrec/models/sequential/base.py b/nextrec/models/sequential/base.py index 1340385..abfd773 100644 --- a/nextrec/models/sequential/base.py +++ b/nextrec/models/sequential/base.py @@ -247,7 +247,7 @@ def compute_loss(self, y_pred: torch.Tensor, y_true: torch.Tensor | None) -> tor flat_labels = flat_labels[valid_mask] if getattr(self, "sequence_loss_mode", "ce") == "sampled_softmax": - return self._sampled_softmax_loss(flat_logits, flat_labels) + return self.sampled_softmax_loss(flat_logits, flat_labels) return F.cross_entropy(flat_logits, flat_labels, reduction="mean") @@ -362,7 +362,7 @@ def evaluate( return {"overall": metrics_dict, "grouped": []} - def _sampled_softmax_loss(self, flat_logits: torch.Tensor, flat_labels: torch.Tensor) -> torch.Tensor: + def sampled_softmax_loss(self, flat_logits: torch.Tensor, flat_labels: torch.Tensor) -> torch.Tensor: """ Approximate full-vocab softmax by sampling a subset of negative items. The positive label is kept for every example; shared negatives are sampled once per batch. diff --git a/nextrec/models/sequential/gru4rec.py b/nextrec/models/sequential/gru4rec.py new file mode 100644 index 0000000..b0d0f82 --- /dev/null +++ b/nextrec/models/sequential/gru4rec.py @@ -0,0 +1,162 @@ +""" +Date: create on 21/04/2026 +Author: Yang Zhou, zyaztec@gmail.com +Reference: +- [1] Tan Y K, Xu X, Liu Y. Improved recurrent neural networks for session-based recommendations. DLRS 2016. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn as nn +from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence + +from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature +from nextrec.basic.layers import EmbeddingLayer +from nextrec.models.sequential.base import BaseSequentialModel +from nextrec.utils.model import select_feature_objects +from nextrec.utils.types import SequenceModeName, TaskTypeInput + + +class GRU4Rec(BaseSequentialModel): + @property + def model_name(self) -> str: + return "GRU4Rec" + + @property + def default_task(self) -> str: + return "sequential" + + def __init__( + self, + sequence_features: list[SequenceFeature], + dense_features: Optional[list[DenseFeature]] = None, + sparse_features: Optional[list[SparseFeature]] = None, + item_history_name: str = "item_history", + hidden_dim: Optional[int] = None, + num_layers: int = 1, + max_seq_len: Optional[int] = None, + dropout_rate: float | None = 0.0, + dropout: float | None = None, + sequence_mode: SequenceModeName = "autoregressive", + target: str | list[str] | None = None, + task: TaskTypeInput | list[TaskTypeInput] | None = None, + embedding_l1_reg: float = 0.0, + dense_l1_reg: float = 0.0, + embedding_l2_reg: float = 0.0, + dense_l2_reg: float = 0.0, + **kwargs, + ): + if not sequence_features: + raise ValueError("[GRU4Rec Error] GRU4Rec requires at least one SequenceFeature.") + if sequence_mode != "autoregressive": + raise ValueError("[GRU4Rec Error] GRU4Rec currently only supports sequence_mode='autoregressive'.") + if num_layers < 1: + raise ValueError(f"[GRU4Rec Error] num_layers must be >= 1, got {num_layers}.") + + if dropout is not None: + dropout_rate = dropout + dropout_rate = float(dropout_rate if dropout_rate is not None else 0.0) + + self.item_history_feature = select_feature_objects( + sequence_features, + [item_history_name], + "item_history_name", + )[0] + self.vocab_size = int(self.item_history_feature.vocab_size) + self.max_seq_len = int(max_seq_len or self.item_history_feature.max_len) + self.padding_idx = ( + self.item_history_feature.padding_idx if self.item_history_feature.padding_idx is not None else 0 + ) + + super().__init__( + dense_features=dense_features, + sparse_features=sparse_features, + sequence_features=sequence_features, + target=target, + task=task or self.default_task, + sequence_mode=sequence_mode, + embedding_l1_reg=embedding_l1_reg, + dense_l1_reg=dense_l1_reg, + embedding_l2_reg=embedding_l2_reg, + dense_l2_reg=dense_l2_reg, + **kwargs, + ) + + self.context_features = [feat for feat in self.all_features if feat.name != self.item_history_feature.name] + self.feature_embedding = EmbeddingLayer(features=self.all_features) + self.item_embedding = self.feature_embedding.embed_dict[self.item_history_feature.embedding_name] + self.context_embedding = self.feature_embedding if self.context_features else None + + item_dim = int(self.item_history_feature.embedding_dim) + context_dim = int(self.feature_embedding.output_dim - item_dim) if self.context_features else 0 + input_dim = item_dim + context_dim + self.hidden_dim = int(hidden_dim or item_dim) + + if input_dim != self.hidden_dim: + self.input_proj = nn.Linear(input_dim, self.hidden_dim) + else: + self.input_proj = nn.Identity() + + gru_dropout = dropout_rate if num_layers > 1 else 0.0 + self.emb_dropout = nn.Dropout(dropout_rate) + self.gru = nn.GRU( + input_size=self.hidden_dim, + hidden_size=self.hidden_dim, + num_layers=num_layers, + batch_first=True, + dropout=gru_dropout, + bias=True, + ) + self.output_proj = nn.Linear(self.hidden_dim, self.vocab_size, bias=False) + + if self.item_embedding.embedding_dim == self.hidden_dim: + self.output_proj.weight = self.item_embedding.weight + + self.register_regularization_weights( + embedding_attr="feature_embedding", + include_modules=["gru", "output_proj", "input_proj"], + ) + + def encode_sequence(self, x) -> tuple[torch.Tensor, torch.Tensor]: + seq, padding_mask, valid_mask, _ = self.prepare_sequence_batch( + x=x, + sequence_name=self.item_history_feature.name, + max_seq_len=self.max_seq_len, + padding_idx=self.padding_idx, + ) + batch_size, seq_len = seq.shape + + item_emb = self.item_embedding(seq) # [B, L, E] + if self.context_features: + context_repr = self.context_embedding(x, self.context_features, squeeze_dim=True) # [B, C] + context_repr = context_repr.unsqueeze(1).expand(-1, seq_len, -1) # [B, L, C] + seq_emb = torch.cat([item_emb, context_repr], dim=-1) # [B, L, E + C] + else: + seq_emb = item_emb + + seq_emb = self.input_proj(seq_emb) + seq_emb = self.emb_dropout(seq_emb) + + valid_lengths = (~padding_mask).sum(dim=1).clamp(min=1) + packed = pack_padded_sequence( + seq_emb, + lengths=valid_lengths.to("cpu"), + batch_first=True, + enforce_sorted=False, + ) + packed_output, _ = self.gru(packed) + hidden_states, _ = pad_packed_sequence( + packed_output, + batch_first=True, + total_length=seq_len, + ) + hidden_states = hidden_states * valid_mask + return hidden_states, padding_mask + + def forward(self, x) -> torch.Tensor: + hidden_states, _ = self.encode_sequence(x) + logits = self.output_proj(hidden_states) # [B, L, V] + return logits diff --git a/pyproject.toml b/pyproject.toml index 047fb9c..cfabdd6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "nextrec" -version = "0.6.12" +version = "0.6.13" description = "A comprehensive recommendation library with match, ranking, and multi-task learning models" readme = "README.md" requires-python = ">=3.10" diff --git a/test/test_generative_models.py b/test/test_generative_models.py index 8507344..6387ca1 100644 --- a/test/test_generative_models.py +++ b/test/test_generative_models.py @@ -14,6 +14,7 @@ import torch from nextrec.basic.features import SequenceFeature +from nextrec.models.sequential.gru4rec import GRU4Rec from nextrec.models.sequential.hstu import ( HSTU, HSTULayer, @@ -341,5 +342,60 @@ def test_hstu_forward_and_loss(self): assert torch.isfinite(loss) +class TestGRU4RecModel: + def test_gru4rec_forward_and_loss(self): + sequence_features = [ + SequenceFeature( + name="item_history", + vocab_size=20, + max_len=6, + embedding_dim=8, + padding_idx=0, + ) + ] + model = GRU4Rec( + sequence_features=sequence_features, + item_history_name="item_history", + hidden_dim=8, + num_layers=1, + max_seq_len=6, + dropout_rate=0.0, + target=["next_item"], + task="generative", + device="cpu", + session_id="gru4rec_test", + ) + model.compile(loss="ce") + + x = { + "item_history": torch.tensor( + [ + [1, 2, 3, 4, 0, 0], + [2, 3, 4, 5, 6, 0], + ], + dtype=torch.long, + ) + } + y_true = torch.tensor( + [ + [2, 3, 4, 5, 0, 0], + [3, 4, 5, 6, 7, 0], + ], + dtype=torch.long, + ) + + y_pred = model.forward(x) + assert y_pred.shape == (2, 6, 20) + assert_no_nan_or_inf(y_pred, "gru4rec_logits") + + loss = model.compute_loss(y_pred, y_true) + assert loss.dim() == 0 + assert torch.isfinite(loss) + + next_item_logits = model.predict_last(x) + assert next_item_logits.shape == (2, 20) + assert_no_nan_or_inf(next_item_logits, "gru4rec_last_logits") + + if __name__ == "__main__": pytest.main([__file__, "-v", "--tb=short"]) diff --git a/tutorials/run_all_sequential_models.py b/tutorials/run_all_sequential_models.py index 5521a54..aba1a57 100644 --- a/tutorials/run_all_sequential_models.py +++ b/tutorials/run_all_sequential_models.py @@ -39,6 +39,7 @@ from __future__ import annotations +from nextrec.models.sequential.gru4rec import GRU4Rec from nextrec.models.sequential.hstu import HSTU from nextrec.models.sequential.sasrec import SASRec from nextrec.models.sequential.tiger import Tiger @@ -166,7 +167,17 @@ def main(): "use_temporal_bias": False, "tie_embeddings": True, }, - ) + ), + ( + GRU4Rec, + "GRU4Rec", + { + "hidden_dim": 16, + "num_layers": 1, + "max_seq_len": 20, + "dropout_rate": 0.1, + }, + ), ] successful = 0