torch-rechub 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

torch_rechub/models/generative/hllm.py

@@ -0,0 +1,249 @@
+"""HLLM: Hierarchical Large Language Model for Recommendation."""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from torch_rechub.utils.hstu_utils import RelPosBias
+
+
+class HLLMTransformerBlock(nn.Module):
+    """Single HLLM Transformer block with self-attention and FFN.
+    
+    This block is similar to HSTULayer but designed for HLLM which uses
+    pre-computed item embeddings as input instead of learnable token embeddings.
+    
+    Args:
+        d_model (int): Hidden dimension.
+        n_heads (int): Number of attention heads.
+        dropout (float): Dropout rate.
+    """
+
+    def __init__(self, d_model=512, n_heads=8, dropout=0.1):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+
+        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
+        self.head_dim = d_model // n_heads
+        self.scale = self.head_dim**-0.5
+
+        # Multi-head self-attention
+        self.W_Q = nn.Linear(d_model, d_model)
+        self.W_K = nn.Linear(d_model, d_model)
+        self.W_V = nn.Linear(d_model, d_model)
+        self.W_O = nn.Linear(d_model, d_model)
+
+        # Feed-forward network
+        ffn_hidden = 4 * d_model
+        self.ffn = nn.Sequential(nn.Linear(d_model, ffn_hidden), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ffn_hidden, d_model), nn.Dropout(dropout))
+
+        # Layer normalization and dropout
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x, rel_pos_bias=None):
+        """Forward pass.
+        
+        Args:
+            x (Tensor): Input of shape (B, L, D).
+            rel_pos_bias (Tensor, optional): Relative position bias.
+            
+        Returns:
+            Tensor: Output of shape (B, L, D).
+        """
+        batch_size, seq_len, _ = x.shape
+
+        # Self-attention with residual
+        residual = x
+        x = self.norm1(x)
+
+        Q = self.W_Q(x)  # (B, L, D)
+        K = self.W_K(x)
+        V = self.W_V(x)
+
+        # Reshape for multi-head attention
+        Q = Q.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        K = K.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+        V = V.view(batch_size, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
+
+        # Attention scores
+        scores = torch.matmul(Q, K.transpose(-2, -1)) * self.scale
+
+        # Causal mask
+        causal_mask = torch.tril(torch.ones(seq_len, seq_len, device=x.device, dtype=torch.bool))
+        scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float('-inf'))
+
+        # Add relative position bias if provided
+        if rel_pos_bias is not None:
+            scores = scores + rel_pos_bias
+
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        attn_output = torch.matmul(attn_weights, V)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, self.d_model)
+        attn_output = self.W_O(attn_output)
+        attn_output = self.dropout(attn_output)
+
+        x = residual + attn_output
+
+        # FFN with residual
+        residual = x
+        x = self.norm2(x)
+        x = self.ffn(x)
+        x = residual + x
+
+        return x
+
+
+class HLLMModel(nn.Module):
+    """HLLM: Hierarchical Large Language Model for Recommendation.
+
+    This is a lightweight implementation of HLLM that uses pre-computed item
+    embeddings as input. The original ByteDance HLLM uses end-to-end training
+    with both Item LLM and User LLM, but this implementation focuses on the
+    User LLM component for resource efficiency.
+
+    Architecture:
+        - Item Embeddings: Pre-computed using LLM (offline, frozen)
+          Format: "{item_prompt}title: {title}description: {description}"
+          where item_prompt = "Compress the following sentence into embedding: "
+        - User LLM: Transformer blocks that model user sequences (trainable)
+        - Scoring Head: Dot product between user representation and item embeddings
+
+    Reference:
+        ByteDance HLLM: https://github.com/bytedance/HLLM
+
+    Args:
+        item_embeddings (Tensor or str): Pre-computed item embeddings of shape
+            (vocab_size, d_model), or path to a .pt file containing embeddings.
+            Generated using the last token's hidden state from an LLM.
+        vocab_size (int): Vocabulary size (number of items).
+        d_model (int): Hidden dimension. Should match item embedding dimension.
+            Default: 512. TinyLlama uses 2048, Baichuan2 uses 4096.
+        n_heads (int): Number of attention heads. Default: 8.
+        n_layers (int): Number of transformer blocks. Default: 4.
+        max_seq_len (int): Maximum sequence length. Default: 256.
+            Official uses MAX_ITEM_LIST_LENGTH=50.
+        dropout (float): Dropout rate. Default: 0.1.
+        use_rel_pos_bias (bool): Whether to use relative position bias. Default: True.
+        use_time_embedding (bool): Whether to use time embeddings. Default: True.
+        num_time_buckets (int): Number of time buckets. Default: 2048.
+        time_bucket_fn (str): Time bucketization function ('sqrt' or 'log'). Default: 'sqrt'.
+        temperature (float): Temperature for NCE scoring. Default: 1.0.
+            Official uses logit_scale = log(1/0.07) ≈ 2.66.
+    """
+
+    def __init__(self, item_embeddings, vocab_size, d_model=512, n_heads=8, n_layers=4, max_seq_len=256, dropout=0.1, use_rel_pos_bias=True, use_time_embedding=True, num_time_buckets=2048, time_bucket_fn='sqrt', temperature=1.0):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.max_seq_len = max_seq_len
+        self.use_time_embedding = use_time_embedding
+        self.num_time_buckets = num_time_buckets
+        self.time_bucket_fn = time_bucket_fn
+        self.temperature = temperature
+
+        # Load item embeddings
+        if isinstance(item_embeddings, str):
+            item_embeddings = torch.load(item_embeddings)
+
+        # Register as buffer (not trainable)
+        self.register_buffer('item_embeddings', item_embeddings.float())
+
+        # Positional embedding
+        self.position_embedding = nn.Embedding(max_seq_len, d_model)
+
+        # Time embedding
+        if use_time_embedding:
+            self.time_embedding = nn.Embedding(num_time_buckets + 1, d_model, padding_idx=0)
+
+        # Transformer blocks
+        self.transformer_blocks = nn.ModuleList([HLLMTransformerBlock(d_model=d_model, n_heads=n_heads, dropout=dropout) for _ in range(n_layers)])
+
+        # Relative position bias
+        self.use_rel_pos_bias = use_rel_pos_bias
+        if use_rel_pos_bias:
+            self.rel_pos_bias = RelPosBias(n_heads, max_seq_len)
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize model parameters."""
+        for name, param in self.named_parameters():
+            if 'weight' in name and len(param.shape) > 1:
+                nn.init.xavier_uniform_(param)
+            elif 'bias' in name:
+                nn.init.constant_(param, 0)
+
+    def _time_diff_to_bucket(self, time_diffs):
+        """Map time differences to bucket indices."""
+        time_diffs = time_diffs.float() / 60.0  # seconds to minutes
+        time_diffs = torch.clamp(time_diffs, min=1e-6)
+
+        if self.time_bucket_fn == 'sqrt':
+            buckets = torch.sqrt(time_diffs).long()
+        elif self.time_bucket_fn == 'log':
+            buckets = torch.log(time_diffs).long()
+        else:
+            raise ValueError(f"Unsupported time_bucket_fn: {self.time_bucket_fn}")
+
+        buckets = torch.clamp(buckets, min=0, max=self.num_time_buckets - 1)
+        return buckets
+
+    def forward(self, seq_tokens, time_diffs=None):
+        """Forward pass.
+        
+        Args:
+            seq_tokens (Tensor): Item token IDs of shape (B, L).
+            time_diffs (Tensor, optional): Time differences in seconds of shape (B, L).
+            
+        Returns:
+            Tensor: Logits of shape (B, L, vocab_size).
+        """
+        batch_size, seq_len = seq_tokens.shape
+
+        # Look up item embeddings
+        item_emb = self.item_embeddings[seq_tokens]  # (B, L, D)
+
+        # Add positional embedding
+        positions = torch.arange(seq_len, dtype=torch.long, device=seq_tokens.device)
+        pos_emb = self.position_embedding(positions)  # (L, D)
+        embeddings = item_emb + pos_emb.unsqueeze(0)  # (B, L, D)
+
+        # Add time embedding if provided
+        if self.use_time_embedding:
+            if time_diffs is None:
+                time_diffs = torch.zeros(batch_size, seq_len, dtype=torch.long, device=seq_tokens.device)
+
+            time_buckets = self._time_diff_to_bucket(time_diffs)
+            time_emb = self.time_embedding(time_buckets)  # (B, L, D)
+            embeddings = embeddings + time_emb
+
+        embeddings = self.dropout(embeddings)
+
+        # Get relative position bias
+        rel_pos_bias = None
+        if self.use_rel_pos_bias:
+            rel_pos_bias = self.rel_pos_bias(seq_len)
+
+        # Pass through transformer blocks
+        x = embeddings
+        for block in self.transformer_blocks:
+            x = block(x, rel_pos_bias=rel_pos_bias)
+
+        # Scoring head: compute dot product with item embeddings
+        # x: (B, L, D), item_embeddings: (V, D)
+        logits = torch.matmul(x, self.item_embeddings.t()) / self.temperature  # (B, L, V)
+
+        return logits

torch_rechub/models/generative/hstu.py

@@ -0,0 +1,189 @@
+"""HSTU: Hierarchical Sequential Transduction Units Model."""
+
+import math
+
+import torch
+import torch.nn as nn
+
+from torch_rechub.basic.layers import HSTUBlock
+from torch_rechub.utils.hstu_utils import RelPosBias
+
+
+class HSTUModel(nn.Module):
+    """HSTU: Hierarchical Sequential Transduction Units model.
+
+    Autoregressive generative recommendation model for sequential data.
+    This module stacks multiple ``HSTUBlock`` layers to capture long-range
+    dependencies in user interaction sequences and predicts the next item.
+
+    Args:
+        vocab_size (int): Vocabulary size (number of distinct items, including PAD).
+        d_model (int): Hidden dimension of the model. Default: 512.
+        n_heads (int): Number of attention heads. Default: 8.
+        n_layers (int): Number of stacked HSTU layers. Default: 4.
+        dqk (int): Dimension of query/key vectors per head. Default: 64.
+        dv (int): Dimension of value vectors per head. Default: 64.
+        max_seq_len (int): Maximum sequence length. Default: 256.
+        dropout (float): Dropout rate applied in the model. Default: 0.1.
+        use_rel_pos_bias (bool): Whether to use relative position bias. Default: True.
+        use_time_embedding (bool): Whether to use time-difference embeddings. Default: True.
+        num_time_buckets (int): Number of time buckets for time embeddings. Default: 2048.
+        time_bucket_fn (str): Function used to bucketize time differences, ``"sqrt"``
+            or ``"log"``. Default: ``"sqrt"``.
+
+    Shape:
+        - Input: ``x`` of shape ``(batch_size, seq_len)``; optional ``time_diffs``
+          of shape ``(batch_size, seq_len)`` representing time differences in seconds.
+        - Output: Logits of shape ``(batch_size, seq_len, vocab_size)``.
+
+    Example:
+        >>> model = HSTUModel(vocab_size=100000, d_model=512)
+        >>> x = torch.randint(0, 100000, (32, 256))
+        >>> time_diffs = torch.randint(0, 86400, (32, 256))
+        >>> logits = model(x, time_diffs)
+        >>> logits.shape
+        torch.Size([32, 256, 100000])
+    """
+
+    def __init__(self, vocab_size, d_model=512, n_heads=8, n_layers=4, dqk=64, dv=64, max_seq_len=256, dropout=0.1, use_rel_pos_bias=True, use_time_embedding=True, num_time_buckets=2048, time_bucket_fn='sqrt'):
+        super().__init__()
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        self.max_seq_len = max_seq_len
+        self.use_time_embedding = use_time_embedding
+        self.num_time_buckets = num_time_buckets
+        self.time_bucket_fn = time_bucket_fn
+
+        # Alpha scaling factor (following the Meta reference implementation)
+        self.alpha = math.sqrt(d_model)
+
+        # Token embedding
+        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
+
+        # Absolute positional embedding
+        self.position_embedding = nn.Embedding(max_seq_len, d_model)
+
+        # Time embedding
+        if use_time_embedding:
+            # Embedding table for time-difference buckets
+            # num_time_buckets + 1: extra bucket reserved for padding
+            self.time_embedding = nn.Embedding(num_time_buckets + 1, d_model, padding_idx=0)
+
+        # HSTU block
+        self.hstu_block = HSTUBlock(d_model=d_model, n_heads=n_heads, n_layers=n_layers, dqk=dqk, dv=dv, dropout=dropout, use_rel_pos_bias=use_rel_pos_bias)
+
+        # Relative position bias
+        self.use_rel_pos_bias = use_rel_pos_bias
+        if use_rel_pos_bias:
+            self.rel_pos_bias = RelPosBias(n_heads, max_seq_len)
+
+        # 输出投影层
+        self.output_projection = nn.Linear(d_model, vocab_size)
+
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+
+        # Initialize model parameters
+        self._init_weights()
+
+    def _init_weights(self):
+        """Initialize model parameters."""
+        for name, param in self.named_parameters():
+            if 'weight' in name and len(param.shape) > 1:
+                nn.init.xavier_uniform_(param)
+            elif 'bias' in name:
+                nn.init.constant_(param, 0)
+
+    def _time_diff_to_bucket(self, time_diffs):
+        """Map raw time differences (in seconds) to discrete bucket indices.
+
+        Following the Meta HSTU implementation, continuous time differences are
+        first converted to minutes and then bucketized using either a square-root
+        or logarithmic transform.
+
+        Args:
+            time_diffs (Tensor): Time differences in seconds,
+                shape ``(batch_size, seq_len)``.
+
+        Returns:
+            Tensor: Integer bucket indices of shape ``(batch_size, seq_len)``.
+        """
+        # Convert seconds to minutes (as in the Meta reference implementation)
+        time_bucket_increments = 60.0
+        time_diffs = time_diffs.float() / time_bucket_increments
+
+        # Ensure non-negative values and avoid log(0)
+        time_diffs = torch.clamp(time_diffs, min=1e-6)
+
+        if self.time_bucket_fn == 'sqrt':
+            # Use the square-root transform: suitable when time differences
+            # are relatively evenly distributed.
+            buckets = torch.sqrt(time_diffs).long()
+        elif self.time_bucket_fn == 'log':
+            # Use the logarithmic transform: suitable when time differences
+            # span several orders of magnitude.
+            buckets = torch.log(time_diffs).long()
+        else:
+            raise ValueError(f"Unsupported time_bucket_fn: {self.time_bucket_fn}")
+
+        # Clamp bucket indices to the valid range [0, num_time_buckets - 1]
+        buckets = torch.clamp(buckets, min=0, max=self.num_time_buckets - 1)
+
+        return buckets
+
+    def forward(self, x, time_diffs=None):
+        """Forward pass.
+
+        Args:
+            x (Tensor): Input token ids of shape ``(batch_size, seq_len)``.
+            time_diffs (Tensor, optional): Time differences in seconds,
+                shape ``(batch_size, seq_len)``. If ``None`` and
+                ``use_time_embedding=True``, all-zero time differences are used.
+
+        Returns:
+            Tensor: Logits over the vocabulary of shape
+                ``(batch_size, seq_len, vocab_size)``.
+        """
+        batch_size, seq_len = x.shape
+
+        # Token embedding with alpha scaling (as in the Meta implementation)
+        token_emb = self.token_embedding(x) * self.alpha  # (B, L, D)
+
+        # Absolute positional embedding
+        positions = torch.arange(seq_len, dtype=torch.long, device=x.device)
+        pos_emb = self.position_embedding(positions)  # (L, D)
+
+        # Combine token and position embeddings
+        embeddings = token_emb + pos_emb.unsqueeze(0)  # (B, L, D)
+
+        # Optional time-difference embedding
+        if self.use_time_embedding:
+            if time_diffs is None:
+                # Fallback: use all-zero time differences when none are provided
+                time_diffs = torch.zeros(batch_size, seq_len, dtype=torch.long, device=x.device)
+
+            # Map raw time differences to bucket indices
+            time_buckets = self._time_diff_to_bucket(time_diffs)  # (B, L)
+
+            # Look up time embeddings and add to the sequence representation
+            time_emb = self.time_embedding(time_buckets)  # (B, L, D)
+
+            # embeddings = token_emb + pos_emb + time_emb
+            embeddings = embeddings + time_emb
+
+        embeddings = self.dropout(embeddings)
+
+        # Relative position bias for self-attention
+        rel_pos_bias = None
+        if self.use_rel_pos_bias:
+            rel_pos_bias = self.rel_pos_bias(seq_len)  # (1, H, L, L)
+
+        # HSTU block
+        hstu_output = self.hstu_block(embeddings, rel_pos_bias=rel_pos_bias)  # (B, L, D)
+
+        # Final projection to vocabulary logits
+        logits = self.output_projection(hstu_output)  # (B, L, V)
+
+        return logits

torch_rechub/models/matching/__init__.py

@@ -1,11 +1,13 @@
-from .dssm import DSSM
-from .youtube_dnn import YoutubeDNN
-from .youtube_sbc import YoutubeSBC
-from .dssm_facebook import FaceBookDSSM
-from .gru4rec import GRU4Rec
-from .comirec import ComirecSA, ComirecDR
-from .mind import MIND
-from .narm import NARM
-from .stamp import STAMP
-from .sasrec import SASRec
-from .sine import SINE
+__all__ = ['DSSM', 'FaceBookDSSM', 'YoutubeDNN', 'YoutubeSBC', 'MIND', 'GRU4Rec', 'NARM', 'SASRec', 'SINE', 'STAMP', 'ComirecDR', 'ComirecSA']
+
+from .comirec import ComirecDR, ComirecSA
+from .dssm import DSSM
+from .dssm_facebook import FaceBookDSSM
+from .gru4rec import GRU4Rec
+from .mind import MIND
+from .narm import NARM
+from .sasrec import SASRec
+from .sine import SINE
+from .stamp import STAMP
+from .youtube_dnn import YoutubeDNN
+from .youtube_sbc import YoutubeSBC