nextrec 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

nextrec/__version__.py

@@ -1 +1 @@
-__version__ = "0.3.1"
+__version__ = "0.3.2"

nextrec/basic/loggers.py

@@ -106,7 +106,7 @@ def setup_logger(session_id: str | os.PathLike | None = None):
 
     console_format = '%(message)s'
     file_format = '%(asctime)s - %(levelname)s - %(message)s'
-    date_format = '%H:%M:%S'
+    date_format = '%Y-%m-%d %H:%M:%S'
     
     logger = logging.getLogger()
     logger.setLevel(logging.INFO)

nextrec/basic/model.py

@@ -216,10 +216,15 @@ class BaseModel(FeatureSpecMixin, nn.Module):
         return train_loader, valid_split
 
     def compile(
-        self, optimizer="adam", optimizer_params: dict | None = None,
-        scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None, scheduler_params: dict | None = None,
-        loss: str | nn.Module | list[str | nn.Module] | None = "bce", loss_params: dict | list[dict] | None = None,
-        loss_weights: int | float | list[int | float] | None = None,):
+        self,
+        optimizer: str | torch.optim.Optimizer = "adam",
+        optimizer_params: dict | None = None,
+        scheduler: str | torch.optim.lr_scheduler._LRScheduler | torch.optim.lr_scheduler.LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | type[torch.optim.lr_scheduler.LRScheduler] | None = None,
+        scheduler_params: dict | None = None,
+        loss: str | nn.Module | list[str | nn.Module] | None = "bce",
+        loss_params: dict | list[dict] | None = None,
+        loss_weights: int | float | list[int | float] | None = None,
+    ):
         optimizer_params = optimizer_params or {}
         self._optimizer_name = optimizer if isinstance(optimizer, str) else optimizer.__class__.__name__
         self._optimizer_params = optimizer_params
@@ -1081,6 +1086,7 @@ class BaseModel(FeatureSpecMixin, nn.Module):
         logger.info(f"  Early Stop Patience:   {self._early_stop_patience}")
         logger.info(f"  Max Gradient Norm:     {self._max_gradient_norm}")
         logger.info(f"  Session ID:            {self.session_id}")
+        logger.info(f"  Features Config Path:  {self.features_config_path}")
         logger.info(f"  Latest Checkpoint:     {self.checkpoint_path}")
         
         logger.info("")
@@ -1195,7 +1201,7 @@ class BaseMatchModel(BaseModel):
     def compile(self, 
                 optimizer: str | torch.optim.Optimizer = "adam",
                 optimizer_params: dict | None = None,
-                scheduler: str | torch.optim.lr_scheduler._LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | None = None,
+                scheduler: str | torch.optim.lr_scheduler._LRScheduler | torch.optim.lr_scheduler.LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | type[torch.optim.lr_scheduler.LRScheduler] | None = None,
                 scheduler_params: dict | None = None,
                 loss: str | nn.Module | list[str | nn.Module] | None = "bce",
                 loss_params: dict | list[dict] | None = None):

nextrec/models/generative/__init__.py

@@ -0,0 +1,5 @@
+from .hstu import HSTU
+
+__all__ = [
+    "HSTU",
+]

nextrec/models/generative/hstu.py

@@ -0,0 +1,399 @@
+"""
+[Info: this version is not released yet, i need to more research on source code and paper]
+Date: create on 01/12/2025
+Checkpoint: edit on 01/12/2025 
+Author: Yang Zhou, zyaztec@gmail.com
+Reference:
+[1] Meta AI. Generative Recommenders (HSTU encoder) — https://github.com/meta-recsys/generative-recommenders
+[2] Ma W, Li P, Chen C, et al. Actions speak louder than words: Trillion-parameter sequential transducers for generative recommendations. arXiv:2402.17152.
+
+Hierarchical Sequential Transduction Unit (HSTU) is the core encoder behind 
+Meta’s Generative Recommenders. It replaces softmax attention with lightweight 
+pointwise activations, enabling extremely deep stacks on long behavior sequences.
+
+In each HSTU layer:
+  (1) Tokens are projected into four streams U, V, Q, K via a shared feed-forward block
+  (2) Softmax-free interactions combine QK^T with Relative Attention Bias (RAB) to encode distance
+  (3) Aggregated context is modulated by U-gating and mapped back through an output projection
+
+Stacking layers yields an efficient causal encoder for next-item 
+generation. With a tied-embedding LM head, HSTU forms 
+a full generative recommendation model.
+
+Key Advantages:
+- Softmax-free attention scales better on deep/long sequences
+- RAB captures temporal structure without extra attention heads
+- Causal masking and padding-aware normalization fit real logs
+- Weight tying reduces parameters and stabilizes training
+- Serves as a drop-in backbone for generative recommendation
+
+HSTU（层次化序列转导单元）是 Meta 生成式推荐的核心编码器，
+用点式激活替代 softmax 注意力，可在长序列上轻松堆叠深层结构。
+
+单层 HSTU 的主要步骤：
+  (1) 将输入一次性映射到 U、V、Q、K 四条通路
+  (2) 利用不含 softmax 的 QK^T 结合相对位置偏置（RAB）建模距离信息
+  (3) 用 U 对聚合上下文进行门控，再映射回输出空间
+
+多层堆叠后，可得到高效的因果编码器；与绑权 LM 头配合即可完成 next-item 预测。
+
+主要优势：
+- 摆脱 softmax，在长序列、深层模型上更易扩展
+- 相对位置偏置稳健刻画时序结构
+- 因果 mask 与 padding 感知归一化贴合真实日志
+- 绑权输出头降低参数量并提升训练稳定性
+- 直接作为生成式推荐的骨干网络
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from nextrec.basic.model import BaseModel
+from nextrec.basic.features import DenseFeature, SequenceFeature, SparseFeature
+
+
+def _relative_position_bucket(
+    relative_position: torch.Tensor,
+    num_buckets: int = 32,
+    max_distance: int = 128,
+) -> torch.Tensor:
+    """
+    map the relative position (i-j) to a bucket in [0, num_buckets).
+    """
+    # only need the negative part for causal attention
+    n = -relative_position
+    n = torch.clamp(n, min=0)
+
+    # when the distance is small, keep it exact
+    max_exact = num_buckets // 2
+    is_small = n < max_exact
+
+    # when the distance is too far, do log scaling
+    large_val = max_exact + ((torch.log(n.float() / max_exact + 1e-6) / math.log(max_distance / max_exact)) * (num_buckets - max_exact)).long()
+    large_val = torch.clamp(large_val, max=num_buckets - 1)
+
+    buckets = torch.where(is_small, n.long(), large_val)
+    return buckets
+
+
+class RelativePositionBias(nn.Module):
+    """
+    Compute relative position bias (RAB) for HSTU attention.
+    The input is the sequence length T, output is [1, num_heads, seq_len, seq_len].
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        num_buckets: int = 32,
+        max_distance: int = 128,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+        self.embedding = nn.Embedding(num_buckets, num_heads)
+
+    def forward(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        # positions: [T]
+        ctx = torch.arange(seq_len, device=device)[:, None]
+        mem = torch.arange(seq_len, device=device)[None, :]
+        rel_pos = mem - ctx  # a matrix to describe all relative positions for each [i,j] pair, shape = [seq_len, seq_len]
+        buckets = _relative_position_bucket(rel_pos, num_buckets=self.num_buckets, max_distance=self.max_distance,)  # map to buckets
+        values = self.embedding(buckets) # embedding vector for each [i,j] pair, shape = [seq_len, seq_len, embedding_dim=num_heads]
+        return values.permute(2, 0, 1).unsqueeze(0) # [1, num_heads, seq_len, seq_len]
+
+class HSTUPointwiseAttention(nn.Module):
+    """
+    Pointwise aggregation attention that implements HSTU without softmax:
+        1) [U, V, Q, K] = split( φ1(f1(X)) ), U: gate, V: value, Q: query, K: key
+        2) AV = φ2(QK^T + rab) V / N, av is attention-weighted value
+        3) Y  = f2( Norm(AV) ⊙ U ), y is output
+    φ1, φ2 use SiLU; Norm uses LayerNorm.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.1,
+        alpha: float | None = None
+    ):
+        super().__init__()
+        if d_model % num_heads != 0:
+            raise ValueError(f"[HSTUPointwiseAttention Error] d_model({d_model}) % num_heads({num_heads}) != 0")
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_head = d_model // num_heads
+        self.alpha = alpha if alpha is not None else (self.d_head ** -0.5)
+        # project input to 4 * d_model for U, V, Q, K
+        self.in_proj = nn.Linear(d_model, 4 * d_model, bias=True)
+        # project output back to d_model
+        self.out_proj = nn.Linear(d_model, d_model, bias=True)
+        self.dropout = nn.Dropout(dropout)
+        self.norm = nn.LayerNorm(d_model)
+
+    def _reshape_heads(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        [B, T, D] -> [B, H, T, d_head]
+        """
+        B, T, D = x.shape
+        return x.view(B, T, self.num_heads, self.d_head).transpose(1, 2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,      # [T, T] with 0 or -inf
+        key_padding_mask: Optional[torch.Tensor] = None,  # [B, T], True = pad
+        rab: Optional[torch.Tensor] = None,            # [1, H, T, T] or None
+    ) -> torch.Tensor:
+        B, T, D = x.shape
+
+        # Eq.(1): [U, V, Q, K] = split( φ1(f1(X)) )
+        h = F.silu(self.in_proj(x))  # [B, T, 4D]
+        U, V, Q, K = h.chunk(4, dim=-1)  # each [B, T, D]
+
+        Qh = self._reshape_heads(Q)  # [B, H, T, d_head]
+        Kh = self._reshape_heads(K)  # [B, H, T, d_head]
+        Vh = self._reshape_heads(V)  # [B, H, T, d_head]
+        Uh = self._reshape_heads(U)  # [B, H, T, d_head]
+
+        # attention logits: QK^T (without 1/sqrt(d) and softmax)
+        logits = torch.matmul(Qh, Kh.transpose(-2, -1)) * self.alpha  # [B, H, T, T]
+
+        # add relative position bias (rab^p), and future extensible rab^t
+        if rab is not None:
+            # rab: [1, H, T, T] or [B, H, T, T]
+            logits = logits + rab
+
+        # construct an "allowed" mask to calculate N
+        # 1 indicates that the (query i, key j) pair is a valid attention pair; 0 indicates it is masked out
+        allowed = torch.ones_like(logits, dtype=torch.float)  # [B, H, T, T]
+
+        # causal mask: attn_mask is usually an upper triangular matrix of -inf with shape [T, T]
+        if attn_mask is not None:
+            allowed = allowed * (attn_mask.view(1, 1, T, T) == 0).float()
+            logits = logits + attn_mask.view(1, 1, T, T)
+
+        # padding mask: key_padding_mask is usually [B, T], True = pad
+        if key_padding_mask is not None:
+            # valid: 1 for non-pad, 0 for pad
+            valid = (~key_padding_mask).float()          # [B, T]
+            valid = valid.view(B, 1, 1, T)               # [B, 1, 1, T]
+            allowed = allowed * valid
+            logits = logits.masked_fill(valid == 0, float("-inf"))
+
+        # Eq.(2): A(X)V(X) = φ2(QK^T + rab) V(X) / N
+        attn = F.silu(logits)  # [B, H, T, T]
+        denom = allowed.sum(dim=-1, keepdim=True)  # [B, H, T, 1]
+        denom = denom.clamp(min=1.0)
+
+        attn = attn / denom  # [B, H, T, T]
+        AV = torch.matmul(attn, Vh)  # [B, H, T, d_head]
+        AV = AV.transpose(1, 2).contiguous().view(B, T, D) # reshape back to [B, T, D]
+        U_flat = Uh.transpose(1, 2).contiguous().view(B, T, D)
+        y = self.out_proj(self.dropout(self.norm(AV) * U_flat))  # [B, T, D]
+        return y
+
+
+class HSTULayer(nn.Module):
+    """
+    HSTUPointwiseAttention + Residual Connection
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        num_heads: int,
+        dropout: float = 0.1,
+        use_rab_pos: bool = True,
+        rab_num_buckets: int = 32,
+        rab_max_distance: int = 128,
+    ):
+        super().__init__()
+        self.attn = HSTUPointwiseAttention(d_model=d_model, num_heads=num_heads, dropout=dropout)
+        self.dropout = nn.Dropout(dropout)
+        self.use_rab_pos = use_rab_pos
+        self.rel_pos_bias = (RelativePositionBias(num_heads=num_heads, num_buckets=rab_num_buckets, max_distance=rab_max_distance) if use_rab_pos else None)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attn_mask: Optional[torch.Tensor] = None,
+        key_padding_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        x: [B, T, D]
+        """
+        B, T, D = x.shape
+        device = x.device
+        rab = None
+        if self.use_rab_pos:
+            rab = self.rel_pos_bias(seq_len=T, device=device) # [1, H, T, T]
+        out = self.attn(x=x, attn_mask=attn_mask, key_padding_mask=key_padding_mask, rab=rab)
+        return x + self.dropout(out)
+
+
+class HSTU(BaseModel):
+    """
+    HSTU encoder for next-item prediction in a causal, generative setup.
+    Pipeline:
+      1) Embed tokens + positions from the behavior history
+      2) Apply stacked HSTU layers with causal mask and optional RAB
+      3) Use the last valid position to produce next-item logits via tied LM head
+    """
+
+    @property
+    def model_name(self) -> str:
+        return "HSTU"
+
+    @property
+    def task_type(self) -> str:
+        return "multiclass"
+
+    def __init__(
+        self,
+        sequence_features: list[SequenceFeature],
+        dense_features: Optional[list[DenseFeature]] = None,
+        sparse_features: Optional[list[SparseFeature]] = None,
+        d_model: Optional[int] = None,
+        num_heads: int = 8,
+        num_layers: int = 4,
+        max_seq_len: int = 200,
+        dropout: float = 0.1,
+        # RAB settings
+        use_rab_pos: bool = True,
+        rab_num_buckets: int = 32,
+        rab_max_distance: int = 128,
+
+        tie_embeddings: bool = True,
+        target: Optional[list[str] | str] = None,
+        optimizer: str = "adam",
+        optimizer_params: Optional[dict] = None,
+        scheduler: Optional[str] = None,
+        scheduler_params: Optional[dict] = None,
+        loss_params: Optional[dict] = 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,
+        device: str = "cpu",
+        **kwargs,
+    ):
+        if not sequence_features:
+            raise ValueError("[HSTU Error] HSTU requires at least one SequenceFeature (user behavior history).")
+
+        # demo version: use the first SequenceFeature as the main sequence
+        self.history_feature = sequence_features[0]
+
+        hidden_dim = d_model or max(int(getattr(self.history_feature, "embedding_dim", 0) or 0), 32)
+        # Make hidden_dim divisible by num_heads
+        if hidden_dim % num_heads != 0:
+            hidden_dim = num_heads * math.ceil(hidden_dim / num_heads)
+
+        self.padding_idx = self.history_feature.padding_idx if self.history_feature.padding_idx is not None else 0
+        self.vocab_size = self.history_feature.vocab_size
+        self.max_seq_len = max_seq_len
+
+        super().__init__(
+            dense_features=dense_features,
+            sparse_features=sparse_features,
+            sequence_features=sequence_features,
+            target=target,
+            task=self.task_type,
+            device=device,
+            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,
+        )
+
+        # token & position embedding (paper usually includes pos embedding / RAB in encoder)
+        self.token_embedding = nn.Embedding(
+            num_embeddings=self.vocab_size,
+            embedding_dim=hidden_dim,
+            padding_idx=self.padding_idx,
+        )
+        self.position_embedding = nn.Embedding(max_seq_len, hidden_dim)
+        self.input_dropout = nn.Dropout(dropout)
+
+        # HSTU layers
+        self.layers = nn.ModuleList([HSTULayer(d_model=hidden_dim, num_heads=num_heads, dropout=dropout, use_rab_pos=use_rab_pos, 
+                                               rab_num_buckets=rab_num_buckets, rab_max_distance=rab_max_distance) for _ in range(num_layers)])
+
+        self.final_norm = nn.LayerNorm(hidden_dim)
+        self.lm_head = nn.Linear(hidden_dim, self.vocab_size, bias=False)
+        if tie_embeddings:
+            self.lm_head.weight = self.token_embedding.weight
+
+        # causal mask buffer
+        self.register_buffer("causal_mask", torch.empty(0), persistent=False)
+        self.ignore_index = self.padding_idx if self.padding_idx is not None else -100
+
+        optimizer_params = optimizer_params or {}
+        scheduler_params = scheduler_params or {}
+        loss_params = loss_params or {}
+        loss_params.setdefault("ignore_index", self.ignore_index)
+
+        self.compile(optimizer=optimizer, optimizer_params=optimizer_params, scheduler=scheduler, scheduler_params=scheduler_params, loss="crossentropy", loss_params=loss_params)
+        self._register_regularization_weights(embedding_attr="token_embedding", include_modules=["layers", "lm_head"])
+
+    def _build_causal_mask(self, seq_len: int, device: torch.device) -> torch.Tensor:
+        """
+        build causal mask of shape [T, T]: upper triangle is -inf, others are 0.
+        This will be added to the logits to simulate causal structure.
+        """
+        if self.causal_mask.numel() == 0 or self.causal_mask.size(0) < seq_len:
+            mask = torch.full((seq_len, seq_len), float("-inf"), device=device)
+            mask = torch.triu(mask, diagonal=1) 
+            self.causal_mask = mask
+        return self.causal_mask[:seq_len, :seq_len]
+
+    def _trim_sequence(self, seq: torch.Tensor) -> torch.Tensor:
+        if seq.size(1) <= self.max_seq_len:
+            return seq
+        return seq[:, -self.max_seq_len :]
+
+    def forward(self, x: dict[str, torch.Tensor]) -> torch.Tensor:
+        seq = x[self.history_feature.name].long()  # [B, T_raw]
+        seq = self._trim_sequence(seq)            # [B, T]
+
+        B, T = seq.shape
+        device = seq.device
+        # position ids: [B, T]
+        pos_ids = torch.arange(T, device=device).unsqueeze(0).expand(B, -1)
+        token_emb = self.token_embedding(seq)         # [B, T, D]
+        pos_emb = self.position_embedding(pos_ids)    # [B, T, D]
+        hidden_states = self.input_dropout(token_emb + pos_emb)
+
+        # padding mask：True = pad
+        padding_mask = seq.eq(self.padding_idx)       # [B, T]
+        attn_mask = self._build_causal_mask(seq_len=T, device=device)  # [T, T]
+
+        for layer in self.layers:
+            hidden_states = layer(x=hidden_states, attn_mask=attn_mask, key_padding_mask=padding_mask)
+        hidden_states = self.final_norm(hidden_states)  # [B, T, D]
+
+        valid_lengths = (~padding_mask).sum(dim=1)  # [B]
+        last_index = (valid_lengths - 1).clamp(min=0)
+        last_hidden = hidden_states[torch.arange(B, device=device), last_index]  # [B, D]
+
+        logits = self.lm_head(last_hidden)  # [B, vocab_size]
+        return logits
+
+    def compute_loss(self, y_pred, y_true):
+        """
+        y_true: [B] or [B, 1], the id of the next item.
+        """
+        if y_true is None:
+            raise ValueError("[HSTU-compute_loss] Training requires y_true (next item id).")
+        labels = y_true.view(-1).long()
+        return self.loss_fn[0](y_pred, labels)

nextrec/utils/optimizer.py

@@ -10,7 +10,7 @@ from typing import Iterable
 
 
 def get_optimizer(
-    optimizer: str = "adam",
+    optimizer: str | torch.optim.Optimizer = "adam",
     params: Iterable[torch.nn.Parameter] | None = None,
     **optimizer_params
 ):
@@ -51,7 +51,11 @@ def get_optimizer(
     return optimizer_fn
 
 
-def get_scheduler(scheduler, optimizer, **scheduler_params):
+def get_scheduler(
+    scheduler: str | torch.optim.lr_scheduler._LRScheduler | torch.optim.lr_scheduler.LRScheduler | type[torch.optim.lr_scheduler._LRScheduler] | type[torch.optim.lr_scheduler.LRScheduler] | None,
+    optimizer,
+    **scheduler_params
+):
     """
     Get learning rate scheduler function.
     
@@ -66,7 +70,7 @@ def get_scheduler(scheduler, optimizer, **scheduler_params):
             scheduler_fn = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, **scheduler_params)
         else:
             raise NotImplementedError(f"Unsupported scheduler: {scheduler}")
-    elif isinstance(scheduler, torch.optim.lr_scheduler._LRScheduler):
+    elif isinstance(scheduler, (torch.optim.lr_scheduler._LRScheduler, torch.optim.lr_scheduler.LRScheduler)):
         scheduler_fn = scheduler
     else:
         raise TypeError(f"Invalid scheduler type: {type(scheduler)}")

{nextrec-0.3.1.dist-info → nextrec-0.3.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nextrec
-Version: 0.3.1
+Version: 0.3.2
 Summary: A comprehensive recommendation library with match, ranking, and multi-task learning models
 Project-URL: Homepage, https://github.com/zerolovesea/NextRec
 Project-URL: Repository, https://github.com/zerolovesea/NextRec
@@ -63,7 +63,7 @@ Description-Content-Type: text/markdown
 ![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.3.1-orange.svg)
+![Version](https://img.shields.io/badge/Version-0.3.2-orange.svg)
 
 English | [中文文档](README_zh.md)
 
@@ -75,13 +75,19 @@ English | [中文文档](README_zh.md)
 
 NextRec is a modern recommendation framework built on PyTorch, delivering a unified experience for modeling, training, and evaluation. It follows a modular design with rich model implementations, data-processing utilities, and engineering-ready training components. NextRec focuses on large-scale industrial recall scenarios on Spark clusters, training on massive offline parquet features.
 
-### Why NextRec
+## Why NextRec
 
 - **Unified feature engineering & data pipeline**: Dense/Sparse/Sequence feature definitions, persistent DataProcessor, and batch-optimized RecDataLoader, matching offline feature training/inference in industrial big-data settings.
 - **Multi-scenario coverage**: Ranking (CTR/CVR), retrieval, multi-task learning, and more marketing/rec models, with a continuously expanding model zoo.
 - **Developer-friendly experience**: Stream processing/training/inference for csv/parquet/pathlike data, plus GPU/MPS acceleration and visualization support.
 - **Efficient training & evaluation**: Standardized engine with optimizers, LR schedulers, early stopping, checkpoints, and detailed logging out of the box.
 
+## Architecture
+
+NextRec adopts a modular and low-coupling engineering design, enabling full-pipeline reusability and scalability across data processing → model construction → training & evaluation → inference & deployment. Its core components include: a Feature-Spec-driven Embedding architecture, the BaseModel abstraction, a set of independent reusable Layers, a unified DataLoader for both training and inference, and a ready-to-use Model Zoo.
+
+![NextRec Architecture](asserts/nextrec_diagram_en.png)
+
 > The project borrows ideas from excellent open-source rec libraries. Early layers referenced [torch-rechub](https://github.com/datawhalechina/torch-rechub) but have been replaced with in-house implementations. torch-rechub remains mature in architecture and models; the author contributed a bit there—feel free to check it out.
 
 ---
@@ -104,7 +110,7 @@ To dive deeper, Jupyter notebooks are available:
 - [Hands on the NextRec framework](/tutorials/notebooks/en/Hands%20on%20nextrec.ipynb)
 - [Using the data processor for preprocessing](/tutorials/notebooks/en/Hands%20on%20dataprocessor.ipynb)
 
-> Current version [0.3.1]: the matching module is not fully polished yet and may have compatibility issues or unexpected errors. Please raise an issue if you run into problems.
+> Current version [0.3.2]: the matching module is not fully polished yet and may have compatibility issues or unexpected errors. Please raise an issue if you run into problems.
 
 ## 5-Minute Quick Start
 

{nextrec-0.3.1.dist-info → nextrec-0.3.2.dist-info}/RECORD

@@ -1,13 +1,13 @@
 nextrec/__init__.py,sha256=CvocnY2uBp0cjNkhrT6ogw0q2bN9s1GNp754FLO-7lo,1117
-nextrec/__version__.py,sha256=r4xAFihOf72W9TD-lpMi6ntWSTKTP2SlzKP1ytkjRbI,22
+nextrec/__version__.py,sha256=vNiWJ14r_cw5t_7UDqDQIVZvladKFGyHH2avsLpN7Vg,22
 nextrec/basic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nextrec/basic/activation.py,sha256=1qs9pq4hT3BUxIiYdYs57axMCm4-JyOBFQ6x7xkHTwM,2849
 nextrec/basic/callback.py,sha256=wwh0I2kKYyywCB-sG9eQXShlpXFJIo75qApJmnI5p6c,1036
 nextrec/basic/features.py,sha256=JtB63jqOIL7zZ5zoTgvEM4fEoqexMz0SMTmowTURk1I,4626
 nextrec/basic/layers.py,sha256=zIa8QsPkOOovjrMAUC94SfhSVTS4R_CXySBr5KAk6i4,24686
-nextrec/basic/loggers.py,sha256=TtTN5NIH8yqY27R2jXxQxfsTIA8XUBPJakx6Bl2ofhI,3724
+nextrec/basic/loggers.py,sha256=VNed0LagpoPSUl2itW8hHT-BSqJHTlQY5pVxIVmm6AE,3733
 nextrec/basic/metrics.py,sha256=YFOaUexHJncc6sPbw2LF2sBnFp-3PLMrjR3aQbBDpGs,20891
-nextrec/basic/model.py,sha256=X1eH9XAxIQla-hVGKUxqEm7QyZucp_tIbx6FWYTa24M,73140
+nextrec/basic/model.py,sha256=Doq5KOYrUHavpSa8RkHbT98ZhbFGRpRsA_9K1A5gU9c,73453
 nextrec/basic/session.py,sha256=oaATn-nzbJ9A6SGbMut9xLV_NSh9_1KmVDeNauS06Ps,4767
 nextrec/data/__init__.py,sha256=COaTyiARV7hEQTT3e74uyCBGmHFQ9rhe6g6Shc-Ualw,1064
 nextrec/data/data_utils.py,sha256=H-isIrs2FPyLSTe7IiFUkn6SQKfO0BkGKmj43C9yLGY,7602
@@ -18,7 +18,8 @@ nextrec/loss/listwise.py,sha256=gxDbO1td5IeS28jKzdE35o1KAYBRdCYoMzyZzfNLhc0,5689
 nextrec/loss/loss_utils.py,sha256=uZ4m9ChLr-UgIc5Yxm1LjwXDDepApQ-Fas8njweZ9qg,2641
 nextrec/loss/pairwise.py,sha256=MN_3Pk6Nj8KCkmUqGT5cmyx1_nQa3TIx_kxXT_HB58c,3396
 nextrec/loss/pointwise.py,sha256=shgdRJwTV7vAnVxHSffOJU4TPQeKyrwudQ8y-R10nYM,7144
-nextrec/models/generative/hstu.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+nextrec/models/generative/__init__.py,sha256=vo8-DloD74cKc1moSH-4GYG99w8Yi8YPGPxh8XDJPoc,50
+nextrec/models/generative/hstu.py,sha256=qTS05XQBjgC5K34A07DSgIITMs1-ADZ8KVb-HEyNh9w,16369
 nextrec/models/generative/tiger.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nextrec/models/match/__init__.py,sha256=ASZB5abqKPhDbk8NErNNNa0DHuWpsVxvUtyEn5XMx6Y,215
 nextrec/models/match/dssm.py,sha256=e0hUqNLJVwTRVz4F4EiO8KLOOprKRBDtI4ID6Y1Tc60,8232
@@ -49,8 +50,8 @@ nextrec/utils/__init__.py,sha256=A3mH6M-DmDBWQ1stIIaTsNzvUy_AKaUWtRmrzU5R3FE,429
 nextrec/utils/common.py,sha256=YTlJkFCvIH5ExiOvg5pNPdRLUQ-h60BX4xTliaXKDsE,1217
 nextrec/utils/embedding.py,sha256=yxYSdFx0cJITh3Gf-K4SdhwRtKGcI0jOsyBgZ0NLa_c,465
 nextrec/utils/initializer.py,sha256=ffYOs5QuIns_d_-5e40iNtg6s1ftgREJN-ueq_NbDQE,1647
-nextrec/utils/optimizer.py,sha256=85ifoy2IQgjPHOqLqr1ho7XBGE_0ry1yEB9efS6C2lM,2446
-nextrec-0.3.1.dist-info/METADATA,sha256=bYvcXVXbnD8hAW8Y-cVKj--ngfd1kCo36atLZFhszT8,15808
-nextrec-0.3.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-nextrec-0.3.1.dist-info/licenses/LICENSE,sha256=2fQfVKeafywkni7MYHyClC6RGGC3laLTXCNBx-ubtp0,1064
-nextrec-0.3.1.dist-info/RECORD,,
+nextrec/utils/optimizer.py,sha256=EUjAGFPeyou_Cv-_2HRvjzut8y_qpAQudc8L2T0k8zw,2706
+nextrec-0.3.2.dist-info/METADATA,sha256=4RFzGjoOmLQUS1wIyJ6edrJgJNZkH9wcxOrQxSLln4w,16319
+nextrec-0.3.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+nextrec-0.3.2.dist-info/licenses/LICENSE,sha256=2fQfVKeafywkni7MYHyClC6RGGC3laLTXCNBx-ubtp0,1064
+nextrec-0.3.2.dist-info/RECORD,,