diff --git a/README.md b/README.md index e0b2d7ba..f8fafe83 100644 --- a/README.md +++ b/README.md @@ -153,6 +153,7 @@ The table below lists the recommendation models/algorithms featured in Cornac. E | 2025 | [Generating Long Semantic IDs in Parallel for Recommendation (RPG)](cornac/models/rpg), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.rpg.recom_rpg), [paper](https://arxiv.org/abs/2506.05781) | Next-Item / Content-Based | [requirements](cornac/models/rpg/requirements.txt), CPU / GPU | [quick-start](examples/rpg_example.py) | 2024 | [Comparative Aspects and Opinions Ranking for Recommendation Explanations (Companion)](cornac/models/companion), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.companion.recom_companion), [paper](https://lthoang.com/assets/publications/mlj24.pdf) | Hybrid / Sentiment / Explainable | CPU | [quick-start](examples/companion_example.py) | | [Hypergraphs with Attention on Reviews (HypAR)](cornac/models/hypar), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.hypar.recom_hypar), [paper](https://doi.org/10.1007/978-3-031-56027-9_14)| Hybrid / Sentiment / Explainable | [requirements](cornac/models/hypar/requirements_cu118.txt), CPU / GPU | [quick-start](https://github.com/PreferredAI/HypAR) +| | [Learnable Item Tokenization for Generative Recommendation (LETTER)](cornac/models/letter), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.letter.recom_letter), [paper](https://arxiv.org/abs/2405.07314) | Next-Item / Content-Based | [requirements](cornac/models/letter/requirements.txt), CPU / GPU | [quick-start](examples/letter_example.py) | 2023 | [Recommender Systems with Generative Retrieval (TIGER)](cornac/models/tiger), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.tiger.recom_tiger), [paper](https://arxiv.org/pdf/2305.05065.pdf) | Next-Item / Content-Based | [requirements](cornac/models/tiger/requirements.txt), CPU / GPU | [quick-start](examples/tiger_example.py) | | [Scalable Approximate NonSymmetric Autoencoder (SANSA)](cornac/models/sansa), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.sansa.recom_sansa), [paper](https://dl.acm.org/doi/10.1145/3604915.3608827) | Collaborative Filtering | [requirements](cornac/models/sansa/requirements.txt), CPU | [quick-start](examples/sansa_movielens.py), [150k-items](examples/sansa_tradesy.py) | 2022 | [Disentangled Multimodal Representation Learning for Recommendation (DMRL)](cornac/models/dmrl), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.dmrl.recom_dmrl), [paper](https://arxiv.org/pdf/2203.05406.pdf) | Content-Based / Text & Image | [requirements](cornac/models/dmrl/requirements.txt), CPU / GPU | [quick-start](examples/dmrl_example.py) diff --git a/cornac/models/__init__.py b/cornac/models/__init__.py index c354d3fa..c4f9f03a 100644 --- a/cornac/models/__init__.py +++ b/cornac/models/__init__.py @@ -58,6 +58,7 @@ from .ibpr import IBPR from .knn import ItemKNN from .knn import UserKNN +from .letter import LETTER from .lightgcn import LightGCN from .lrppm import LRPPM from .mcf import MCF diff --git a/cornac/models/letter/README.md b/cornac/models/letter/README.md new file mode 100644 index 00000000..ea72b59d --- /dev/null +++ b/cornac/models/letter/README.md @@ -0,0 +1,96 @@ +# LETTER + +Cornac implementation of **LETTER** (Learnable Item Tokenization for Generative Recommendation, Wang et al., CIKM 2024, [arXiv:2405.07314](https://arxiv.org/abs/2405.07314)). LETTER adds collaborative alignment and code-assignment diversity to an RQ-VAE tokenizer, then trains a T5 generator to predict the next item's four-token Semantic ID. This implementation follows the [released code](https://github.com/HonghuiBao2000/LETTER) for both stages. + +## Requirements + +Install the optional PyTorch, Transformers, and constrained-k-means dependencies: + +```bash +pip install -r cornac/models/letter/requirements.txt +``` + +LETTER needs two aligned feature matrices covering every item known to the train, validation, and test splits: + +- item content embeddings, supplied through `FeatureModality`; and +- 32-dimensional collaborative item embeddings and their raw item IDs, + supplied through `cf_embeddings` and `cf_embedding_ids` (the paper uses + SASRec item embeddings). LETTER remaps these rows to Cornac's global item + indices during fitting. + +## Usage + +```python +from cornac.data import FeatureModality +from cornac.eval_methods import NextItemEvaluation +from cornac.models import LETTER +from cornac.models.letter import LETTER_BEAUTY_CONFIG + +eval_method = NextItemEvaluation.from_splits( + train_data=train, + val_data=val, + test_data=test, + mode="last", + item_feature=FeatureModality(features=item_embeddings, ids=item_ids), +) + +model = LETTER( + **{ + **LETTER_BEAUTY_CONFIG, + "cf_embeddings": sasrec_item_embeddings_32d, + "cf_embedding_ids": sasrec_item_ids, + "device": "auto", + "seed": 42, + } +) +``` + +See [`examples/letter_example.py`](../../../examples/letter_example.py) for a small two-stage API example with stand-in features. `LETTER_BEAUTY_CONFIG` is the reproduction recipe; `LETTER_CONFIG` keeps the paper-wide recommended regularization weights. + +## Training + +LETTER is trained in two stages. The tokenizer is a four-level RQ-VAE with collaborative alignment, code-assignment diversity, and collision handling. The generator is a T5 model that predicts the four-token Semantic ID of the next item. `LETTER_BEAUTY_CONFIG` contains the released Beauty training settings. + +The released “ranking-guided” objective uses a temperature of 1.0, making it equivalent to ordinary token cross-entropy. Generator AdamW weight decay follows Hugging Face Trainer: layer-normalization and bias parameters are placed in a zero-decay group. + +`precomputed_semantic_ids` is a positional table whose rows must already follow Cornac's global item-index order. If semantic IDs come with raw item IDs, remap them through the evaluation method's `global_iid_map` before constructing `LETTER`. + +## Beauty reproduction + +All results use the Amazon Beauty 2014 5-core interactions, seed 42, and a chronological leave-last-out split. The Cornac results have two scopes: + +- **Generator + author IDs** skips tokenizer training and evaluates the Cornac generator with the authors' Semantic IDs over all 12,101 released items. +- **End-to-end** trains both stages using Sentence-T5-base title+description embeddings and locally trained 32-dimensional SASRec embeddings. It covers the 12,068-item universe for which both inputs are available. + +The complete [Beauty generator example](beauty_example.py) consumes the authors' released `Beauty.index.json` and `Beauty.inter.json`, then trains once and evaluates multiple beam widths from the same checkpoint: + +```bash +python -m cornac.models.letter.beauty_example Beauty.index.json \ + --interaction-file Beauty.inter.json --beams 20 50 +``` + +Omit `--interaction-file` to use Cornac's Amazon loader with an ASIN-keyed Semantic-ID file. For a numeric-keyed index in that mode, pass a numeric-to-ASIN JSON mapping through `--item-id-map`. + +### Results + +Recall is abbreviated as R and NDCG as N. + +| System | Beams | R@5 | N@5 | R@10 | N@10 | +| ----------------------------- | ----: | -----: | -----: | -----: | -----: | +| LETTER paper | 20 | 0.0431 | 0.0286 | 0.0672 | 0.0364 | +| Released code + author IDs | 20 | 0.0413 | 0.0268 | 0.0645 | 0.0343 | +| Cornac generator + author IDs | 20 | 0.0420 | 0.0279 | 0.0656 | 0.0354 | +| Cornac generator + author IDs | 50 | 0.0429 | 0.0282 | 0.0670 | 0.0360 | +| Cornac end-to-end | 50 | 0.0356 | 0.0241 | 0.0560 | 0.0307 | + +### Interpretation + +Across the four reported metrics, it is within `0.0016` of the paper and `0.0011` of the released-code rerun, supporting generator fidelity. + +The end-to-end run trained the 10,000-epoch LETTER tokenizer and generator with reproducible substitute inputs, reducing Semantic-ID collisions from 74 to 2. It used 50 beams and predates the corrected generator optimizer, so it is retained as evidence of the full local training pipeline rather than as a controlled generator comparison. + +The authors' Beauty index contains 12,101 items but only 12,088 unique Semantic IDs. The released evaluator deduplicates these candidate strings, whereas Cornac retains the raw items and gives items sharing an ID the same score; metrics involving those collisions are therefore not directly comparable. + +### Reproduction limitation + +The released repository provides the final Beauty Semantic-ID table, but not the trained tokenizer checkpoint, content embeddings, or SASRec checkpoint used to produce it. Exact reproduction of the paper's learned Semantic IDs therefore requires those missing artifacts; the end-to-end result above is a runnable substitute rather than an exact reconstruction of the paper's tokenizer inputs. diff --git a/cornac/models/letter/__init__.py b/cornac/models/letter/__init__.py new file mode 100644 index 00000000..725f2597 --- /dev/null +++ b/cornac/models/letter/__init__.py @@ -0,0 +1,17 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# 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 .letter_config import LETTER_BEAUTY_CONFIG, LETTER_CONFIG +from .recom_letter import LETTER diff --git a/cornac/models/letter/beauty_example.py b/cornac/models/letter/beauty_example.py new file mode 100644 index 00000000..2d81175d --- /dev/null +++ b/cornac/models/letter/beauty_example.py @@ -0,0 +1,163 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# 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. +# ============================================================================ +"""Reproduce LETTER's Beauty generator evaluation with author Semantic IDs. + +Pair the authors' numeric-keyed ``Beauty.index.json`` with their released +``Beauty.inter.json``. Alternatively, use an ASIN-keyed Semantic-ID file with +Cornac's Amazon loader. Tokenizer reproduction is out of scope because its +Beauty inputs are not public. +""" + +import argparse +import json +import re + +import numpy as np + +from cornac.data import FeatureModality +from cornac.datasets import amazon_review +from cornac.eval_methods import NextItemEvaluation +from cornac.metrics import NDCG, Recall +from cornac.models import LETTER +from cornac.models.letter import LETTER_BEAUTY_CONFIG + + +def load_semantic_ids(path, item_id_map_path=None): + """Load four-level Semantic IDs keyed by raw item ID.""" + with open(path) as stream: + raw_ids = json.load(stream) + if item_id_map_path is None: + item_id_map = None + else: + with open(item_id_map_path) as stream: + item_id_map = json.load(stream) + + semantic_ids = {} + for source_id, tokens in raw_ids.items(): + if len(tokens) != 4: + raise ValueError( + f"semantic ID for {source_id!r} has {len(tokens)} levels; expected 4" + ) + item_id = source_id if item_id_map is None else item_id_map[str(source_id)] + codes = [] + for token in tokens: + if isinstance(token, int): + codes.append(token) + continue + match = re.fullmatch(r"<[a-d]_(\d+)>", token) + if match is None: + raise ValueError(f"invalid semantic token {token!r}") + codes.append(int(match.group(1))) + if item_id in semantic_ids: + raise ValueError(f"duplicate mapped item ID {item_id!r}") + semantic_ids[item_id] = codes + return semantic_ids + + +def align_semantic_ids(eval_method, semantic_ids): + """Arrange raw-ID-keyed codes in Cornac's global item-index order.""" + missing = set(eval_method.global_iid_map) - set(semantic_ids) + if missing: + raise ValueError(f"semantic-ID file is missing {len(missing)} Beauty items") + aligned = np.empty((eval_method.total_items, 4), dtype="int64") + for item_id, item_index in eval_method.global_iid_map.items(): + aligned[item_index] = semantic_ids[item_id] + return aligned + + +def load_released_feedback(path): + """Convert the authors' ordered Beauty sequences to Cornac feedback.""" + with open(path) as stream: + sequences = json.load(stream) + return [ + (user_id, str(item_id), 1.0, timestamp) + for user_id, item_ids in sequences.items() + for timestamp, item_id in enumerate(item_ids) + ] + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("semantic_id_file", help="JSON item-to-Semantic-ID mapping") + parser.add_argument( + "--item-id-map", + help="JSON numeric-ID-to-ASIN map for the authors' Beauty.index.json", + ) + parser.add_argument( + "--interaction-file", + help="authors' Beauty.inter.json; otherwise use Cornac's Amazon loader", + ) + parser.add_argument( + "--beams", + nargs="+", + type=int, + default=[20, 50], + help="beam widths evaluated from the same trained generator", + ) + parser.add_argument("--device", default="auto") + args = parser.parse_args() + if any(width <= 0 for width in args.beams): + parser.error("--beams values must be positive") + + semantic_ids = load_semantic_ids(args.semantic_id_file, args.item_id_map) + item_ids = list(semantic_ids) + feedback = ( + amazon_review.load_feedback("beauty") + if args.interaction_file is None + else load_released_feedback(args.interaction_file) + ) + eval_method = NextItemEvaluation.leave_last_out( + feedback, + fmt="UIRT", + mode="last", + exclude_unknowns=False, + item_feature=FeatureModality( + features=np.zeros((len(item_ids), 1), dtype="float32"), + ids=item_ids, + ), + verbose=True, + ) + aligned_ids = align_semantic_ids(eval_method, semantic_ids) + + config = dict(LETTER_BEAUTY_CONFIG) + config.update( + precomputed_semantic_ids=aligned_ids, + n_beams=args.beams[0], + device=args.device, + seed=42, + verbose=True, + ) + model = LETTER(name=f"LETTER-b{args.beams[0]}", **config) + metrics = [Recall(k=5), NDCG(k=5), Recall(k=10), NDCG(k=10)] + + results = {} + for index, beam_width in enumerate(dict.fromkeys(args.beams)): + if index: + model.trainable = False + model.n_beams = beam_width + model.name = f"LETTER-b{beam_width}" + test_result, _ = eval_method.evaluate( + model, metrics=metrics, user_based=False, show_validation=False + ) + results[str(beam_width)] = { + key: float(value) + for key, value in test_result.metric_avg_results.items() + if "(s)" not in key + } + print(json.dumps(results, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/cornac/models/letter/letter.py b/cornac/models/letter/letter.py new file mode 100644 index 00000000..085f2402 --- /dev/null +++ b/cornac/models/letter/letter.py @@ -0,0 +1,518 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# 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. +# ============================================================================ +"""Neural modules for LETTER (Wang et al., CIKM 2024). + +The implementation follows the authors' released ``RQ-VAE`` and +``LETTER-TIGER`` code. In particular, the tokenizer uses constrained k-means +initialization, a Sinkhorn assignment on the last residual codebook, and the +released collaborative/diversity losses. The generator uses a tied T5 +vocabulary and predicts EOS after the four semantic-ID tokens. +""" + +import random + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def _letter_mlp(input_dim, hidden_dims, output_dim, dropout=0.0): + """MLP used by the released LETTER tokenizer.""" + dims = [input_dim, *hidden_dims, output_dim] + layers = [] + for index, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + layers.extend((nn.Dropout(dropout), nn.Linear(in_dim, out_dim))) + if index != len(dims) - 2: + layers.append(nn.ReLU()) + model = nn.Sequential(*layers) + for module in model.modules(): + if isinstance(module, nn.Linear): + nn.init.xavier_normal_(module.weight) + nn.init.zeros_(module.bias) + return model + + +def _constrained_kmeans(data, n_clusters, initial=False, n_jobs=10): + """Run LETTER's constrained k-means and return centers and labels. + + ``k-means-constrained`` is imported here so importing :mod:`cornac` does + not require the optional tokenizer dependency. + """ + try: + from joblib import parallel_backend + from k_means_constrained import KMeansConstrained + except ImportError as exc: + raise ImportError( + "LETTER requires k-means-constrained. Install the dependencies in " + "cornac/models/letter/requirements.txt." + ) from exc + + n_samples = len(data) + if n_samples < n_clusters: + raise ValueError( + "constrained codebook initialization needs at least as many items " + f"as codes ({n_samples} < {n_clusters})" + ) + cap = 50 if initial else 10 + size_min = min(n_samples // (n_clusters * 2), cap) + if size_min < 1: + raise ValueError( + "LETTER's constrained k-means needs at least two samples per cluster" + ) + size_max = min( + size_min * 4 if initial else n_clusters * 6, n_samples - 1 + ) + clusterer = KMeansConstrained( + n_clusters=n_clusters, + size_min=size_min, + size_max=size_max, + max_iter=10, + n_init=10, + n_jobs=n_jobs, + verbose=False, + ) + values = data.detach().cpu().numpy() + # Joblib otherwise memory-maps arrays above 1 MiB as read-only for its + # workers, but k-means-constrained's Cython center update needs a writable + # buffer. Keep the released parallelism while passing normal arrays. + with parallel_backend("loky", max_nbytes=None): + clusterer.fit(values) + centers = torch.as_tensor( + clusterer.cluster_centers_, dtype=data.dtype, device=data.device + ) + labels = torch.as_tensor( + clusterer.labels_, dtype=torch.long, device=data.device + ) + return centers, labels + + +def sinkhorn_assignment(distances, epsilon, n_iters): + """Balanced assignment from the released LETTER Sinkhorn routine.""" + maximum = distances.max() + minimum = distances.min() + middle = (maximum + minimum) / 2 + amplitude = maximum - middle + 1e-5 + centered = ((distances - middle) / amplitude).double() + + q = torch.exp(-centered / epsilon) + batch_size, n_codes = q.shape + q = q / q.sum() + for _ in range(n_iters): + q = q / q.sum(dim=1, keepdim=True) + q = q / batch_size + q = q / q.sum(dim=0, keepdim=True) + q = q / n_codes + return (q * batch_size).argmax(dim=1) + + +class LETTERRQVAE(nn.Module): + """Residual-quantized autoencoder from the official LETTER release.""" + + def __init__( + self, + input_dim, + hidden_dims=(2048, 1024, 512, 256, 128, 64), + latent_dim=32, + num_levels=4, + codebook_size=256, + commitment_weight=0.25, + n_clusters=10, + sk_epsilons=None, + sk_iters=50, + dropout=0.0, + kmeans_n_jobs=10, + ): + super().__init__() + self.num_levels = num_levels + self.codebook_size = codebook_size + self.latent_dim = latent_dim + self.commitment_weight = commitment_weight + self.n_clusters = n_clusters + self.sk_epsilons = tuple( + [0.0] * (num_levels - 1) + [0.003] + if sk_epsilons is None + else sk_epsilons + ) + if len(self.sk_epsilons) != num_levels: + raise ValueError("sk_epsilons must contain one value per codebook") + self.sk_iters = sk_iters + self.kmeans_n_jobs = kmeans_n_jobs + self.encoder = _letter_mlp(input_dim, hidden_dims, latent_dim, dropout) + # nn.Embedding initializes before the official code zeros each + # k-means-initialized table. Consume the same RNG draws so the decoder + # and subsequent shuffled loader start from the released seed state. + codebooks = torch.empty(num_levels, codebook_size, latent_dim) + nn.init.normal_(codebooks) + self.codebooks = nn.Parameter(codebooks.zero_()) + self.decoder = _letter_mlp( + latent_dim, tuple(reversed(hidden_dims)), input_dim, dropout + ) + self._div_labels = None + + @staticmethod + def _distances(x, codebook): + return ( + x.square().sum(dim=1, keepdim=True) + + codebook.square().sum(dim=1).unsqueeze(0) + - 2 * x @ codebook.t() + ) + + def _assign(self, distances, level, use_sinkhorn): + epsilon = self.sk_epsilons[level] + if use_sinkhorn and epsilon > 0: + return sinkhorn_assignment(distances, epsilon, self.sk_iters) + return distances.argmin(dim=1) + + @torch.no_grad() + def initialize_codebooks(self, x): + """Initialize every residual level on the complete item collection.""" + residual = self.encoder(x) + for level in range(self.num_levels): + centers, _ = _constrained_kmeans( + residual, + self.codebook_size, + initial=True, + n_jobs=self.kmeans_n_jobs, + ) + self.codebooks[level].copy_(centers) + distances = self._distances(residual, centers) + ids = self._assign(distances, level, use_sinkhorn=True) + residual = residual - centers[ids] + + @torch.no_grad() + def update_diversity_clusters(self): + """Refresh the ten constrained codebook groups once per epoch.""" + self._div_labels = [ + _constrained_kmeans( + self.codebooks[level], + self.n_clusters, + initial=False, + n_jobs=self.kmeans_n_jobs, + )[1] + for level in range(self.num_levels) + ] + + def _diversity_loss(self, codebook, selected, ids, level): + if self._div_labels is None: + return codebook.new_zeros(()) + labels = self._div_labels[level].tolist() + groups = { + group: [index for index, label in enumerate(labels) if label == group] + for group in range(self.n_clusters) + } + valid_rows = [] + positives = [] + for row, code in enumerate(ids.tolist()): + choices = groups[labels[code]] + # The released 256-code/10-cluster constraints guarantee siblings. + # Keep reduced smoke-test codebooks finite when that is impossible. + if len(choices) < 2: + continue + positive = random.choice(choices) + while positive == code: + positive = random.choice(choices) + valid_rows.append(row) + positives.append(positive) + if not valid_rows: + return codebook.new_zeros(()) + valid_rows = torch.as_tensor( + valid_rows, dtype=torch.long, device=ids.device + ) + targets = torch.as_tensor(positives, dtype=torch.long, device=ids.device) + selected_ids = ids[valid_rows] + logits = selected[valid_rows] @ codebook.t() + logits = logits.clone() + logits.scatter_(1, selected_ids[:, None], -1e12) + return F.cross_entropy(logits, targets) + + def _quantize(self, z, use_sinkhorn=True): + all_ids = [] + quantized = torch.zeros_like(z) + residual = z + quant_loss = z.new_zeros(()) + diversity_loss = z.new_zeros(()) + for level in range(self.num_levels): + codebook = self.codebooks[level] + distances = self._distances(residual, codebook) + ids = self._assign(distances, level, use_sinkhorn) + selected = codebook[ids] + diversity_loss = diversity_loss + self._diversity_loss( + codebook, selected, ids, level + ) + quant_loss = quant_loss + F.mse_loss( + selected, residual.detach() + ) + self.commitment_weight * F.mse_loss( + selected.detach(), residual + ) + + # Match VectorQuantizer.forward: straight-through at every level, + # then subtract that value before quantizing the next residual. + selected_st = residual + (selected - residual).detach() + residual = residual - selected_st + quantized = quantized + selected_st + all_ids.append(ids) + scale = float(self.num_levels) + return ( + torch.stack(all_ids, dim=1), + quantized, + quant_loss / scale, + diversity_loss / scale, + ) + + @staticmethod + def _cf_loss(quantized, cf_batch): + labels = torch.arange(quantized.size(0), device=quantized.device) + return F.cross_entropy(quantized @ cf_batch.t(), labels) + + def forward(self, x, cf_batch=None): + """Return IDs, reconstruction, and all released loss components.""" + z = self.encoder(x) + ids, quantized, loss_rq, loss_div = self._quantize( + z, use_sinkhorn=True + ) + reconstruction = self.decoder(quantized) + loss_recon = F.mse_loss(reconstruction, x) + if cf_batch is None: + loss_cf = x.new_zeros(()) + else: + if cf_batch.size(1) != self.latent_dim: + raise ValueError( + "official LETTER requires CF embeddings to match the " + f"{self.latent_dim}-d tokenizer latent (got {cf_batch.size(1)})" + ) + loss_cf = self._cf_loss(quantized, cf_batch) + return ids, reconstruction, loss_recon, loss_rq, loss_cf, loss_div + + @torch.no_grad() + def encode(self, x, use_sinkhorn=False): + ids, _, _, _ = self._quantize( + self.encoder(x), use_sinkhorn=use_sinkhorn + ) + return ids + + @torch.no_grad() + def resolve_collisions(self, x, codes, max_iters=20): + """Reassign colliding groups with last-level Sinkhorn, as released.""" + resolved = codes.clone() + for _ in range(max_iters): + groups = {} + for item, row in enumerate(resolved.tolist()): + groups.setdefault(tuple(row), []).append(item) + collisions = [items for items in groups.values() if len(items) > 1] + if not collisions: + break + for items in collisions: + item_ids = torch.as_tensor(items, device=x.device) + resolved[item_ids] = self.encode( + x[item_ids], use_sinkhorn=True + ) + return resolved + + +class LETTERSeq2Seq(nn.Module): + """T5 generator matching the released LETTER-TIGER parameterization.""" + + eos_token_id = 1 + pad_token_id = 0 + + def __init__( + self, + level_sizes, + d_model=128, + d_ff=1024, + num_heads=6, + d_kv=64, + num_enc_layers=4, + num_dec_layers=4, + dropout=0.1, + base_vocab_size=32100, + temperature=1.0, + code_values=None, + ): + super().__init__() + from transformers import T5Config, T5ForConditionalGeneration + + self.level_sizes = [int(size) for size in level_sizes] + self.num_levels = len(self.level_sizes) + self.temperature = temperature + self.base_vocab_size = base_vocab_size + if code_values is None: + code_values = [range(size) for size in self.level_sizes] + code_values = [list(map(int, values)) for values in code_values] + vocab_size = base_vocab_size + sum(map(len, code_values)) + config = T5Config( + vocab_size=vocab_size, + d_model=d_model, + d_ff=d_ff, + d_kv=d_kv, + num_heads=num_heads, + num_layers=num_enc_layers, + num_decoder_layers=num_dec_layers, + dropout_rate=dropout, + decoder_start_token_id=self.pad_token_id, + pad_token_id=self.pad_token_id, + eos_token_id=self.eos_token_id, + use_cache=False, + ) + self.t5 = T5ForConditionalGeneration(config) + + # T5Tokenizer.add_tokens receives a sorted set in the official code. + # Reproduce its lexicographic code-token order within each level. + token_ids = [] + offset = base_vocab_size + for level, (size, values) in enumerate( + zip(self.level_sizes, code_values) + ): + prefix = chr(ord("a") + level) + order = sorted(values, key=lambda code: f"<{prefix}_{code}>") + inverse = torch.zeros(size, dtype=torch.long) + for rank, code in enumerate(order): + inverse[code] = offset + rank + token_ids.append(inverse) + offset += len(values) + self.register_buffer("code_token_ids", torch.stack(token_ids)) + + def semantic_tokens(self, codes): + levels = torch.arange(codes.size(-1), device=codes.device) + return self.code_token_ids[levels, codes] + + def forward(self, enc_tokens, enc_mask, target_sids): + target_tokens = self.semantic_tokens(target_sids) + eos = target_tokens.new_full((target_tokens.size(0), 1), self.eos_token_id) + labels = torch.cat((target_tokens, eos), dim=1) + decoder_inputs = torch.cat( + (labels.new_full((labels.size(0), 1), self.pad_token_id), labels[:, :-1]), + dim=1, + ) + output = self.t5( + input_ids=enc_tokens, + attention_mask=enc_mask, + decoder_input_ids=decoder_inputs, + use_cache=False, + return_dict=True, + ) + logits = output.logits / self.temperature + return F.cross_entropy(logits.reshape(-1, logits.size(-1)), labels.reshape(-1)) + + def _decoder_logits(self, decoder_inputs, enc_out, enc_mask): + hidden = self.t5.decoder( + input_ids=decoder_inputs, + encoder_hidden_states=enc_out, + encoder_attention_mask=enc_mask, + use_cache=False, + return_dict=True, + ).last_hidden_state + if self.t5.config.tie_word_embeddings: + hidden = hidden * (self.t5.model_dim**-0.5) + return self.t5.lm_head(hidden[:, -1]) / self.temperature + + @torch.no_grad() + def generate_beam(self, enc_tokens, enc_mask, n_beams, prefix_children): + enc_out = self.t5.encoder( + input_ids=enc_tokens, attention_mask=enc_mask, return_dict=True + ).last_hidden_state + beams = [()] + beam_scores = enc_out.new_zeros(1) + for level, size in enumerate(self.level_sizes): + n_current = len(beams) + decoder_inputs = torch.full( + (n_current, level + 1), + self.pad_token_id, + dtype=torch.long, + device=enc_out.device, + ) + if level: + previous = torch.as_tensor(beams, device=enc_out.device) + decoder_inputs[:, 1:] = self.semantic_tokens(previous) + logits = self._decoder_logits( + decoder_inputs, + enc_out.expand(n_current, -1, -1), + enc_mask.expand(n_current, -1), + ) + full_log_probs = F.log_softmax(logits, dim=-1) + code_ids = self.code_token_ids[level] + log_probs = full_log_probs.index_select(1, code_ids) + allowed = torch.full_like(log_probs, float("-inf")) + for row, beam in enumerate(beams): + allowed[row, prefix_children[level][beam]] = 0.0 + totals = (beam_scores[:, None] + log_probs + allowed).flatten() + width = min(n_beams, int(torch.isfinite(totals).sum())) + top = totals.topk(width) + beams = [ + beams[index // size] + (index % size,) + for index in top.indices.tolist() + ] + beam_scores = top.values + + previous = torch.as_tensor(beams, device=enc_out.device) + decoder_inputs = torch.cat( + ( + previous.new_full((len(beams), 1), self.pad_token_id), + self.semantic_tokens(previous), + ), + dim=1, + ) + eos_logits = self._decoder_logits( + decoder_inputs, + enc_out.expand(len(beams), -1, -1), + enc_mask.expand(len(beams), -1), + ) + beam_scores = beam_scores + F.log_softmax(eos_logits, dim=-1)[ + :, self.eos_token_id + ] + order = beam_scores.argsort(descending=True) + return [beams[index] for index in order.tolist()], beam_scores[order].cpu().numpy() + + @torch.no_grad() + def score_all_items(self, enc_tokens, enc_mask, sid_table, batch_size): + enc_out = self.t5.encoder( + input_ids=enc_tokens, attention_mask=enc_mask, return_dict=True + ).last_hidden_state + scores = enc_out.new_empty(sid_table.size(0)) + for start in range(0, sid_table.size(0), batch_size): + codes = sid_table[start : start + batch_size] + n_items = codes.size(0) + target_tokens = self.semantic_tokens(codes) + decoder_inputs = torch.cat( + ( + target_tokens.new_full((n_items, 1), self.pad_token_id), + target_tokens, + ), + dim=1, + ) + hidden = self.t5.decoder( + input_ids=decoder_inputs, + encoder_hidden_states=enc_out.expand(n_items, -1, -1), + encoder_attention_mask=enc_mask.expand(n_items, -1), + use_cache=False, + return_dict=True, + ).last_hidden_state + if self.t5.config.tie_word_embeddings: + hidden = hidden * (self.t5.model_dim**-0.5) + item_scores = hidden.new_zeros(n_items) + labels = torch.cat( + ( + target_tokens, + target_tokens.new_full((n_items, 1), self.eos_token_id), + ), + dim=1, + ) + for position in range(self.num_levels + 1): + logits = self.t5.lm_head(hidden[:, position]) / self.temperature + item_scores += F.log_softmax(logits, dim=-1).gather( + 1, labels[:, position : position + 1] + ).squeeze(1) + scores[start : start + n_items] = item_scores + return scores.cpu().numpy() diff --git a/cornac/models/letter/letter_config.py b/cornac/models/letter/letter_config.py new file mode 100644 index 00000000..512922e9 --- /dev/null +++ b/cornac/models/letter/letter_config.py @@ -0,0 +1,78 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# 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. +# ============================================================================ +"""Released LETTER tokenizer and LETTER-TIGER training configurations.""" + +# Shared settings from RQ-VAE/main.py, LETTER-TIGER/ckpt/TIGER/config.json, +# LETTER-TIGER/utils.py, and LETTER-TIGER/run_train.sh. On one GPU, four +# accumulated 256-example minibatches reproduce the published two-GPU +# effective batch (2 devices x 256 x 2 accumulation = 1024 examples). +LETTER_CONFIG = { + "feature_standardize": False, + "rqvae_num_levels": 4, + "rqvae_codebook_size": 256, + "rqvae_latent_dim": 32, + "rqvae_hidden_dims": (2048, 1024, 512, 256, 128, 64), + "rqvae_beta": 0.25, + "rqvae_quant_loss_weight": 1.0, + "rqvae_sk_epsilon": 0.003, + "rqvae_sk_iters": 50, + "rqvae_kmeans_jobs": 10, + "rqvae_learning_rate": 1e-3, + "rqvae_batch_size": 1024, + "rqvae_weight_decay": 1e-4, + "rqvae_n_epochs": 10000, + "n_clusters": 10, + "collision_resolve_iters": 20, + # Paper-wide recommended regularization values. The released Beauty + # command overrides these below. + "cf_weight": 0.02, + "diversity_weight": 1e-3, + # Released 4+4 T5 and generation recipe. + "d_model": 128, + "d_ff": 1024, + "d_kv": 64, + "num_heads": 6, + "num_enc_layers": 4, + "num_dec_layers": 4, + "dropout": 0.1, + "letter_base_vocab_size": 32100, + "ranking_temperature": 1.0, + "max_len": 20, + "n_epochs": 200, + "learning_rate": 5e-4, + "weight_decay": 0.01, + "batch_size": 256, + "gradient_accumulation_steps": 4, + "lr_schedule": "cosine", + "warmup_ratio": 0.01, + "model_selection": "best", + "val_eval_every": 1, + "val_sample": None, + "val_batch_size": 256, + "early_stopping_patience": 20, + "max_grad_norm": 1.0, + "scoring": "beam", + "n_beams": 20, + "scoring_batch_size": 256, +} + + +# RQ-VAE/tokenize.sh selects the Beauty tokenizer trained for 10,000 epochs +# with alpha=0.1 and beta=1e-4. +LETTER_BEAUTY_CONFIG = { + **LETTER_CONFIG, + "cf_weight": 0.1, + "diversity_weight": 1e-4, +} diff --git a/cornac/models/letter/recom_letter.py b/cornac/models/letter/recom_letter.py new file mode 100644 index 00000000..b32cde15 --- /dev/null +++ b/cornac/models/letter/recom_letter.py @@ -0,0 +1,549 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# 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. +# ============================================================================ + +import math +import random +from collections import defaultdict + +import numpy as np +from tqdm.auto import trange + +from ..tiger.recom_tiger import TIGER + + +def _generator_optimizer_groups(model, weight_decay): + """Match Hugging Face Trainer's AdamW decay groups for the T5 generator.""" + decay = [] + no_decay = [] + for name, parameter in model.named_parameters(): + if not parameter.requires_grad: + continue + target = no_decay if "bias" in name or "layer_norm" in name else decay + target.append(parameter) + return [ + {"params": decay, "weight_decay": weight_decay}, + {"params": no_decay, "weight_decay": 0.0}, + ] + + +class LETTER(TIGER): + """LETTER: LEarnable Tokenizer for generaTivE Recommendation. + + LETTER trains a four-level RQ-VAE tokenizer with collaborative and + code-assignment-diversity regularization, then trains the released + LETTER-TIGER generator over the resulting semantic IDs. The official + collaborative features are 32-dimensional SASRec item embeddings. Their + raw item IDs are used to align the rows with Cornac's global item indices. + + This class accepts TIGER's public parameters plus the LETTER-specific + arguments below. :data:`~cornac.models.letter.LETTER_BEAUTY_CONFIG` + contains the authors' released Beauty recipe. + + Parameters + ---------- + cf_embeddings: array-like or None + Collaborative item embeddings of shape ``(n_items, rqvae_latent_dim)``. + Required when ``cf_weight`` is non-zero. + cf_embedding_ids: array-like or None + Raw item IDs corresponding to the rows of ``cf_embeddings``. Required + whenever ``cf_embeddings`` is provided. Rows are aligned to Cornac's + global item indices during fitting. + cf_weight: float, default: 0.02 + Collaborative loss weight (alpha). + diversity_weight: float, default: 0.001 + Diversity loss weight (beta). + n_clusters: int, default: 10 + Constrained codebook groups used by the diversity loss. + rqvae_quant_loss_weight: float, default: 1.0 + Weight of the complete residual-quantization loss. + rqvae_sk_epsilon: float, default: 0.003 + Sinkhorn epsilon on the final residual codebook. Earlier levels use + nearest-code assignment, matching the release. + rqvae_sk_iters: int, default: 50 + Number of Sinkhorn normalization iterations. + rqvae_kmeans_jobs: int, default: 10 + Worker count used by the released constrained-k-means calls. + collision_resolve_iters: int, default: 20 + Maximum official post-tokenization collision-reassignment passes. + ranking_temperature: float, default: 1.0 + Temperature in the released ranking loss. The published value 1.0 is + ordinary token cross-entropy. + gradient_accumulation_steps: int, default: 1 + Number of generator minibatches per optimizer update. + warmup_ratio: float, default: 0.01 + Fraction of generator optimizer updates used for linear warmup. + early_stopping_patience: int or None, default: 20 + Non-improving epoch validations before stopping. Validation uses the + released token loss and restores the lowest-loss checkpoint. + val_batch_size: int, default: 256 + Batch size for validation loss. + max_grad_norm: float, default: 1.0 + Generator gradient clipping threshold used by Hugging Face Trainer. + letter_base_vocab_size: int, default: 32100 + Size of the base T5 SentencePiece vocabulary before semantic tokens. + precomputed_semantic_ids: array-like or None + Optional integer semantic-ID table of shape ``(n_items, num_levels)``. + When provided, skip tokenizer training and train only the released + LETTER-TIGER generator. + """ + + def __init__( + self, + name="LETTER", + cf_embeddings=None, + cf_embedding_ids=None, + cf_weight=0.02, + diversity_weight=0.001, + n_clusters=10, + rqvae_quant_loss_weight=1.0, + rqvae_sk_epsilon=0.003, + rqvae_sk_iters=50, + rqvae_kmeans_jobs=10, + collision_resolve_iters=20, + ranking_temperature=1.0, + gradient_accumulation_steps=1, + warmup_ratio=0.01, + early_stopping_patience=20, + val_batch_size=256, + max_grad_norm=1.0, + letter_base_vocab_size=32100, + precomputed_semantic_ids=None, + **kwargs, + ): + kwargs["tokenizer"] = "rqvae" + super().__init__(name=name, **kwargs) + if ranking_temperature <= 0: + raise ValueError("ranking_temperature must be positive") + if gradient_accumulation_steps <= 0: + raise ValueError("gradient_accumulation_steps must be positive") + if not 0 <= warmup_ratio <= 1: + raise ValueError("warmup_ratio must be between 0 and 1") + if early_stopping_patience is not None and early_stopping_patience <= 0: + raise ValueError("early_stopping_patience must be positive or None") + if val_batch_size <= 0: + raise ValueError("val_batch_size must be positive") + if rqvae_kmeans_jobs == 0: + raise ValueError("rqvae_kmeans_jobs must be non-zero") + self.cf_embeddings = ( + None + if cf_embeddings is None + else np.asarray(cf_embeddings, dtype="float32") + ) + self.cf_embedding_ids = ( + None if cf_embedding_ids is None else list(cf_embedding_ids) + ) + if (self.cf_embeddings is None) != (self.cf_embedding_ids is None): + raise ValueError( + "cf_embeddings and cf_embedding_ids must be provided together" + ) + if self.cf_embeddings is not None: + if self.cf_embeddings.ndim != 2: + raise ValueError("cf_embeddings must be a 2-dimensional array") + if len(self.cf_embedding_ids) != self.cf_embeddings.shape[0]: + raise ValueError( + f"cf_embedding_ids has {len(self.cf_embedding_ids)} entries " + f"but cf_embeddings has {self.cf_embeddings.shape[0]} rows" + ) + if len(set(self.cf_embedding_ids)) != len(self.cf_embedding_ids): + raise ValueError("cf_embedding_ids must not contain duplicates") + self.cf_weight = cf_weight + self.diversity_weight = diversity_weight + self.n_clusters = n_clusters + self.rqvae_quant_loss_weight = rqvae_quant_loss_weight + self.rqvae_sk_epsilon = rqvae_sk_epsilon + self.rqvae_sk_iters = rqvae_sk_iters + self.rqvae_kmeans_jobs = rqvae_kmeans_jobs + self.collision_resolve_iters = collision_resolve_iters + self.ranking_temperature = ranking_temperature + self.gradient_accumulation_steps = gradient_accumulation_steps + self.warmup_ratio = warmup_ratio + self.early_stopping_patience = early_stopping_patience + self.val_batch_size = val_batch_size + self.max_grad_norm = max_grad_norm + self.letter_base_vocab_size = letter_base_vocab_size + self.precomputed_semantic_ids = ( + None + if precomputed_semantic_ids is None + else np.asarray(precomputed_semantic_ids, dtype="int64") + ) + + def _get_cf_embeddings(self): + if self.cf_weight and self.cf_embeddings is None: + raise ValueError("LETTER requires cf_embeddings when cf_weight is non-zero") + if self.cf_embeddings is None: + return None + if self.cf_embeddings.shape[1] != self.rqvae_latent_dim: + raise ValueError( + "official LETTER uses same-dimensional collaborative and " + f"tokenizer representations; expected {self.rqvae_latent_dim}, " + f"got {self.cf_embeddings.shape[1]}" + ) + + row_by_id = {raw_id: row for row, raw_id in enumerate(self.cf_embedding_ids)} + missing = [raw_id for raw_id in self.iid_map if raw_id not in row_by_id] + if missing: + raise ValueError( + f"cf_embedding_ids is missing {len(missing)} item(s) known to Cornac" + ) + + aligned = np.empty((self.total_items, self.rqvae_latent_dim), dtype="float32") + for raw_id, item_idx in self.iid_map.items(): + aligned[item_idx] = self.cf_embeddings[row_by_id[raw_id]] + return aligned + + def _fit_rqvae(self, torch, feats_t): + from .letter import LETTERRQVAE + + cf_embeddings = self._get_cf_embeddings() + cf_t = ( + None + if cf_embeddings is None + else torch.as_tensor(cf_embeddings, device=self.device_) + ) + + seed = self.seed if self.seed is not None else 0 + random.seed(seed) + np.random.seed(seed) + self.rqvae = LETTERRQVAE( + input_dim=feats_t.size(1), + hidden_dims=self.rqvae_hidden_dims, + latent_dim=self.rqvae_latent_dim, + num_levels=self.rqvae_num_levels, + codebook_size=self.rqvae_codebook_size, + commitment_weight=self.rqvae_beta, + n_clusters=self.n_clusters, + sk_epsilons=[0.0] * (self.rqvae_num_levels - 1) + [self.rqvae_sk_epsilon], + sk_iters=self.rqvae_sk_iters, + kmeans_n_jobs=self.rqvae_kmeans_jobs, + ).to(self.device_) + init_loader = torch.utils.data.DataLoader( + range(feats_t.size(0)), + batch_size=feats_t.size(0), + shuffle=True, + ) + init_ids = next(iter(init_loader)).to(self.device_) + self.rqvae.initialize_codebooks(feats_t[init_ids]) + optimizer = torch.optim.AdamW( + self.rqvae.parameters(), + lr=self.rqvae_learning_rate, + weight_decay=self.rqvae_weight_decay, + ) + + train_loader = torch.utils.data.DataLoader( + range(feats_t.size(0)), + batch_size=self.rqvae_batch_size, + shuffle=True, + ) + progress = trange( + 1, + self.rqvae_n_epochs + 1, + disable=not self.verbose, + desc="LETTER RQ-VAE", + ) + for _ in progress: + self.rqvae.train() + if self.diversity_weight: + self.rqvae.update_diversity_clusters() + total_loss = 0.0 + count = 0 + for item_ids in train_loader: + item_ids = item_ids.to(self.device_) + batch = feats_t[item_ids] + cf_batch = None if cf_t is None else cf_t[item_ids] + _, _, recon, quant, collaborative, diversity = self.rqvae( + batch, cf_batch + ) + loss = ( + recon + + self.rqvae_quant_loss_weight + * (quant + self.diversity_weight * diversity) + + self.cf_weight * collaborative + ) + optimizer.zero_grad() + loss.backward() + optimizer.step() + total_loss += loss.item() * len(batch) + count += len(batch) + progress.set_postfix(loss=total_loss / count) + + def _tokenize(self, torch, feats_t): + if self.precomputed_semantic_ids is not None: + codes = self.precomputed_semantic_ids + expected = (self.total_items, self.rqvae_num_levels) + if codes.shape != expected: + raise ValueError( + f"precomputed_semantic_ids has shape {codes.shape}; expected {expected}" + ) + if codes.size and ( + codes.min() < 0 or codes.max() >= self.rqvae_codebook_size + ): + raise ValueError( + f"precomputed_semantic_ids values must be in [0, {self.rqvae_codebook_size})" + ) + unique = np.unique(codes, axis=0) + self.sid_collisions_before = len(codes) - len(unique) + self.sid_collisions_after = self.sid_collisions_before + self.sid_code_utilization = [ + int(np.unique(codes[:, level]).size) + for level in range(self.rqvae_num_levels) + ] + return codes.copy() + + self._fit_rqvae(torch, feats_t) + self.rqvae.eval() + codes = torch.cat( + [ + self.rqvae.encode( + feats_t[start : start + self.rqvae_batch_size], + use_sinkhorn=False, + ) + for start in range(0, feats_t.size(0), self.rqvae_batch_size) + ] + ) + self.sid_collisions_before = len(codes) - len(torch.unique(codes, dim=0)) + resolved = self.rqvae.resolve_collisions( + feats_t, codes, max_iters=self.collision_resolve_iters + ) + self.sid_collisions_after = len(resolved) - len(torch.unique(resolved, dim=0)) + self.sid_code_utilization = [ + int(resolved[:, level].unique().numel()) + for level in range(self.rqvae_num_levels) + ] + return resolved.cpu().numpy().astype("int64") + + def _build_semantic_ids(self, codes): + """Build the official fixed-length IDs without TIGER's dedup level.""" + self.sid_table = np.asarray(codes, dtype="int64") + self.level_sizes = [self.rqvae_codebook_size] * self.rqvae_num_levels + children = [defaultdict(set) for _ in self.level_sizes] + sid_to_items = defaultdict(list) + for item, row in enumerate(self.sid_table): + sid = tuple(int(value) for value in row) + for level in range(len(sid)): + children[level][sid[:level]].add(sid[level]) + sid_to_items[sid].append(item) + # NumPy advanced indexing lets TIGER's beam scorer assign the same + # token score to every item left in a collision, instead of dropping it. + self.sid_to_item = { + sid: np.asarray(items, dtype="int64") for sid, items in sid_to_items.items() + } + self.prefix_children = [ + { + prefix: np.fromiter(sorted(tokens), dtype="int64") + for prefix, tokens in level.items() + } + for level in children + ] + if self.verbose: + collisions = sum(len(items) - 1 for items in sid_to_items.values()) + print( + f"LETTER semantic IDs: {len(codes)} items, {collisions} unresolved collisions" + ) + + def _training_rows(self): + rows = [] + uir_tuple = self.train_set.uir_tuple + for mapped_ids in self.train_set.sessions.values(): + items = [int(item) for item in uir_tuple[1][mapped_ids]] + for position in range(1, len(items)): + rows.append((items[:position][-self.max_len :], items[position])) + return rows + + def _encoder_batch(self, torch, histories): + sequences = [] + for history in histories: + item_ids = np.asarray(history[-self.max_len :], dtype="int64") + tokens = self.enc_token_table[item_ids].reshape(-1).tolist() + sequences.append(tokens + [self.model.eos_token_id]) + max_tokens = max(len(sequence) for sequence in sequences) + batch = np.full( + (len(sequences), max_tokens), self.model.pad_token_id, dtype="int64" + ) + for row, sequence in enumerate(sequences): + batch[row, : len(sequence)] = sequence + tokens = torch.as_tensor(batch, dtype=torch.long, device=self.device_) + return tokens, (tokens != self.model.pad_token_id).float() + + def _validation_loss(self, torch, val_sessions): + was_training = self.model.training + self.model.eval() + total = 0.0 + count = 0 + with torch.no_grad(): + for start in range(0, len(val_sessions), self.val_batch_size): + batch = val_sessions[start : start + self.val_batch_size] + histories = [items[:-1] for _, items in batch] + targets = [items[-1] for _, items in batch] + enc_tokens, enc_mask = self._encoder_batch(torch, histories) + target_sids = torch.as_tensor( + self.sid_table[targets], dtype=torch.long, device=self.device_ + ) + loss = self.model(enc_tokens, enc_mask, target_sids) + total += loss.item() * len(batch) + count += len(batch) + if was_training: + self.model.train() + return total / count if count else float("inf") + + def _fit_seq2seq(self, torch, val_set): + from .letter import LETTERSeq2Seq + + # Tokenizer and generator are separate programs in the release; both + # restart from the configured seed. + seed = self.seed if self.seed is not None else 0 + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + self.rng = np.random.RandomState(seed) + + self.model = LETTERSeq2Seq( + level_sizes=self.level_sizes, + d_model=self.d_model, + d_ff=self.d_ff, + num_heads=self.num_heads, + d_kv=self.d_kv, + num_enc_layers=self.num_enc_layers, + num_dec_layers=self.num_dec_layers, + dropout=self.dropout, + base_vocab_size=self.letter_base_vocab_size, + temperature=self.ranking_temperature, + code_values=[ + np.unique(self.sid_table[:, level]) + for level in range(self.rqvae_num_levels) + ], + ).to(self.device_) + self.pad_idx = self.total_items + code_table = torch.as_tensor( + self.sid_table, dtype=torch.long, device=self.device_ + ) + self.enc_token_table = self.model.semantic_tokens(code_table).cpu().numpy() + + optimizer = torch.optim.AdamW( + _generator_optimizer_groups(self.model, self.weight_decay), + lr=self.learning_rate, + ) + rows = self._training_rows() + if not rows: + raise ValueError("LETTER needs at least one next-item training prefix") + train_loader = torch.utils.data.DataLoader( + range(len(rows)), batch_size=self.batch_size, shuffle=True + ) + batches_per_epoch = max(1, len(train_loader)) + updates_per_epoch = max( + 1, math.ceil(batches_per_epoch / self.gradient_accumulation_steps) + ) + total_updates = max(1, updates_per_epoch * self.n_epochs) + warmup_updates = math.ceil(total_updates * self.warmup_ratio) + + def lr_lambda(step): + if step < warmup_updates: + return step / max(1, warmup_updates) + progress = (step - warmup_updates) / max(1, total_updates - warmup_updates) + return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) + + scheduler = ( + torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) + if self.lr_schedule == "cosine" + else None + ) + select_best = self.model_selection == "best" and val_set is not None + val_sessions = self._val_sessions(val_set) if select_best else [] + best_state = None + best_loss = float("inf") + non_improving = 0 + + progress = trange( + 1, self.n_epochs + 1, disable=not self.verbose, desc="LETTER-TIGER" + ) + for epoch in progress: + self.current_epoch = epoch + self.model.train() + optimizer.zero_grad() + total_loss = 0.0 + count = 0 + for batch_index, indices in enumerate(train_loader, start=1): + batch = [rows[index] for index in indices.tolist()] + histories = [history for history, _ in batch] + targets = [target for _, target in batch] + enc_tokens, enc_mask = self._encoder_batch(torch, histories) + target_sids = torch.as_tensor( + self.sid_table[targets], dtype=torch.long, device=self.device_ + ) + loss = self.model(enc_tokens, enc_mask, target_sids) + (loss / self.gradient_accumulation_steps).backward() + if ( + batch_index % self.gradient_accumulation_steps == 0 + or batch_index == batches_per_epoch + ): + torch.nn.utils.clip_grad_norm_( + self.model.parameters(), self.max_grad_norm + ) + optimizer.step() + if scheduler is not None: + scheduler.step() + optimizer.zero_grad() + total_loss += loss.item() * len(batch) + count += len(batch) + progress.set_postfix(loss=total_loss / count) + + if select_best and epoch % self.val_eval_every == 0: + val_loss = self._validation_loss(torch, val_sessions) + if val_loss < best_loss: + best_loss = val_loss + best_state = { + name: value.detach().cpu().clone() + for name, value in self.model.state_dict().items() + } + non_improving = 0 + else: + non_improving += 1 + if ( + self.early_stopping_patience is not None + and non_improving >= self.early_stopping_patience + ): + break + + self.best_val_loss = None if best_state is None else best_loss + if best_state is not None: + self.model.load_state_dict(best_state) + + def score(self, user_idx, history_items, **kwargs): + import torch + + if len(history_items) == 0: + return np.ones(self.total_items, dtype="float") + self._ensure_device(torch) + enc_tokens, enc_mask = self._encoder_batch( + torch, [list(history_items)[-self.max_len :]] + ) + self.model.eval() + with torch.no_grad(): + if self.scoring == "beam": + beams, log_probs = self.model.generate_beam( + enc_tokens, enc_mask, self.n_beams, self.prefix_children + ) + scores = np.full(self.total_items, -1e10, dtype="float") + for sid, log_prob in zip(beams, log_probs): + scores[self.sid_to_item[sid]] = log_prob + else: + sid_table = torch.as_tensor( + self.sid_table, dtype=torch.long, device=self.device_ + ) + scores = self.model.score_all_items( + enc_tokens, enc_mask, sid_table, self.scoring_batch_size + ).astype("float") + return scores diff --git a/cornac/models/letter/requirements.txt b/cornac/models/letter/requirements.txt new file mode 100644 index 00000000..31764bd0 --- /dev/null +++ b/cornac/models/letter/requirements.txt @@ -0,0 +1,3 @@ +torch>=1.12.0 +transformers>=4.30.0 +k-means-constrained>=0.7.3 diff --git a/cornac/models/tiger/recom_tiger.py b/cornac/models/tiger/recom_tiger.py index d4a68051..1ec31bbc 100644 --- a/cornac/models/tiger/recom_tiger.py +++ b/cornac/models/tiger/recom_tiger.py @@ -554,7 +554,7 @@ def fit(self, train_set, val_set=None): # keep pickles portable across GPU/CPU boxes; moved back in score() self.model.to("cpu").eval() - if self.tokenizer == "rqvae": + if self.tokenizer == "rqvae" and hasattr(self, "rqvae"): self.rqvae.to("cpu").eval() return self diff --git a/docs/source/api_ref/models.rst b/docs/source/api_ref/models.rst index d8592075..edbebbbc 100644 --- a/docs/source/api_ref/models.rst +++ b/docs/source/api_ref/models.rst @@ -16,6 +16,11 @@ Comparative Aspects and Opinions Ranking for Recommendation Explanations (Compan .. automodule:: cornac.models.companion.recom_companion :members: +Learnable Item Tokenization for Generative Recommendation (LETTER) +------------------------------------------------------------------ +.. automodule:: cornac.models.letter.recom_letter + :members: + Generating Long Semantic IDs in Parallel for Recommendation (RPG) ----------------------------------------------------------------- .. automodule:: cornac.models.rpg.recom_rpg diff --git a/examples/README.md b/examples/README.md index fb4b1330..9924135e 100644 --- a/examples/README.md +++ b/examples/README.md @@ -132,6 +132,8 @@ [tiger_example.py](tiger_example.py) - Generative retrieval with semantic IDs (TIGER) on Amazon Beauty with Sentence-T5 item content embeddings, reproducing the paper's leave-last-out protocol. +[letter_example.py](letter_example.py) - Learnable item tokenization (LETTER): constrained/Sinkhorn RQ-VAE with collaborative + diversity regularization and a Semantic-ID T5 generator, with Diginetica dataset. + [rpg_example.py](rpg_example.py) - Parallel generation of long unordered semantic IDs (RPG): OPQ tokenizer + multi-token prediction + graph-guided decoding, with Diginetica dataset. ---- diff --git a/examples/letter_example.py b/examples/letter_example.py new file mode 100644 index 00000000..9bdb2848 --- /dev/null +++ b/examples/letter_example.py @@ -0,0 +1,118 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# +# 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. +# ============================================================================ +"""LETTER (learnable tokenizer for generative retrieval) on Diginetica. + +LETTER replaces TIGER's RQ-VAE tokenizer with the released LETTER tokenizer: +(1) a collaborative InfoNCE loss aligning each item's semantic ID with a +precomputed collaborative (CF) item embedding, and (2) a diversity loss that +spreads codebook usage. Its downstream generator follows the released tied +T5-vocabulary/EOS objective and epoch-level validation-loss early stopping. + +Two things are precomputed and passed in: + * item CONTENT embeddings -> the evaluation method's FeatureModality, e.g. + with sentence-transformers:: + + from sentence_transformers import SentenceTransformer + content = SentenceTransformer("sentence-t5-base").encode(titles) + + * item COLLABORATIVE embeddings -> ``LETTER(cf_embeddings=..., + cf_embedding_ids=...)``, typically the item embeddings of a trained CF + model (SASRec in the paper). Raw IDs are supplied so LETTER can align the + rows to Cornac's global item indices. + +Diginetica ships without item text/CF vectors in Cornac, so this example uses +random vectors as stand-ins -- replace both with real embeddings for +meaningful semantic IDs. +""" + +import numpy as np +import torch + +import cornac +from cornac.data import FeatureModality +from cornac.datasets import diginetica +from cornac.eval_methods import NextItemEvaluation +from cornac.metrics import MRR, NDCG, Recall +from cornac.models import LETTER, TIGER + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +print(f"using device: {DEVICE}") + +train_data = diginetica.load_train() +val_data = diginetica.load_val() +test_data = diginetica.load_test() +print("data loaded") + +item_ids = sorted({tup[2] for tup in train_data + val_data + test_data}) +rng = np.random.RandomState(123) +print( + "NOTE: using random content + CF features as stand-ins; replace with real " + "content embeddings and trained-CF item embeddings (see module docstring)." +) +content = rng.randn(len(item_ids), 768).astype("float32") +cf_embeddings = rng.randn(len(item_ids), 32).astype("float32") # e.g. SASRec item embs + +next_item_eval = NextItemEvaluation.from_splits( + train_data=train_data, + val_data=val_data, + test_data=test_data, + exclude_unknowns=True, + verbose=True, + fmt="USIT", + item_feature=FeatureModality(features=content, ids=item_ids), +) + +models = [ + LETTER( # lightweight example budget; use LETTER_BEAUTY_CONFIG to reproduce + cf_embeddings=cf_embeddings, + cf_embedding_ids=item_ids, + cf_weight=0.02, + diversity_weight=1e-3, + rqvae_num_levels=4, + rqvae_codebook_size=256, + rqvae_latent_dim=32, + rqvae_n_epochs=200, + n_epochs=50, + batch_size=256, + max_len=20, + scoring="beam", + n_beams=50, + device=DEVICE, + verbose=True, + seed=123, + ), + TIGER( # baseline: same pipeline, plain RQ-VAE tokenizer + rqvae_num_levels=4, + rqvae_codebook_size=256, + rqvae_latent_dim=32, + rqvae_n_epochs=200, + n_epochs=50, + batch_size=256, + max_len=20, + scoring="beam", + n_beams=50, + device=DEVICE, + verbose=True, + seed=123, + ), +] + +metrics = [NDCG(k=10), NDCG(k=50), Recall(k=10), Recall(k=50), MRR()] + +cornac.Experiment( + eval_method=next_item_eval, + models=models, + metrics=metrics, +).run()