From ccd310cd0ecb24b21584ee71d3a6b94a9460440e Mon Sep 17 00:00:00 2001 From: hieuddo Date: Mon, 2 Nov 2026 00:44:03 +0800 Subject: [PATCH 1/2] Add DIGER Semantic ID recommender Implement frequency-gated residual quantization and T5 generation with lazy optional dependencies, Beauty presets, and shared Semantic ID evaluation. Include a feature-based example, CPU regression tests, and verified retained-checkpoint inference results with explicit reproduction limits. --- README.md | 1 + cornac/models/__init__.py | 1 + cornac/models/diger/README.md | 116 ++++++++++ cornac/models/diger/__init__.py | 21 ++ cornac/models/diger/diger.py | 311 +++++++++++++++++++++++++++ cornac/models/diger/diger_config.py | 40 ++++ cornac/models/diger/recom_diger.py | 249 +++++++++++++++++++++ cornac/models/diger/requirements.txt | 2 + cornac/models/diger/training.py | 201 +++++++++++++++++ docs/source/api_ref/models.rst | 5 + examples/README.md | 2 + examples/diger_example.py | 81 +++++++ tests/cornac/models/test_diger.py | 159 ++++++++++++++ 13 files changed, 1189 insertions(+) create mode 100644 cornac/models/diger/README.md create mode 100644 cornac/models/diger/__init__.py create mode 100644 cornac/models/diger/diger.py create mode 100644 cornac/models/diger/diger_config.py create mode 100644 cornac/models/diger/recom_diger.py create mode 100644 cornac/models/diger/requirements.txt create mode 100644 cornac/models/diger/training.py create mode 100644 examples/diger_example.py create mode 100644 tests/cornac/models/test_diger.py diff --git a/README.md b/README.md index f971ae52..910a874a 100644 --- a/README.md +++ b/README.md @@ -150,6 +150,7 @@ The table below lists the recommendation models/algorithms featured in Cornac. E | Year | Model and Paper | Type | Environment | Example | | :--: | --------------- | :--: | :---------: | :-----: | +| 2026 | [Differentiable Semantic ID for Generative Recommendation (DIGER)](cornac/models/diger), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.diger.recom_diger), [paper](https://arxiv.org/abs/2601.19711) | Next-Item / Content-Based | [requirements](cornac/models/diger/requirements.txt), CPU / GPU | [quick-start](examples/diger_example.py) | | 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) | | [Diffusion-based Generative Recommendation Model (DiffGRM)](cornac/models/diffgrm), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.diffgrm.recom_diffgrm), [paper](https://arxiv.org/abs/2510.21805) | Next-Item / Content-Based | [requirements](cornac/models/diffgrm/requirements.txt), CPU / GPU | [quick-start](examples/diffgrm_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) diff --git a/cornac/models/__init__.py b/cornac/models/__init__.py index 994a344d..4c031d5b 100644 --- a/cornac/models/__init__.py +++ b/cornac/models/__init__.py @@ -43,6 +43,7 @@ from .cvae import CVAE from .cvaecf import CVAECF from .dmrl import DMRL +from .diger import DIGER from .diffgrm import DiffGRM from .dnntsp import DNNTSP from .ease import EASE diff --git a/cornac/models/diger/README.md b/cornac/models/diger/README.md new file mode 100644 index 00000000..3d35c7e4 --- /dev/null +++ b/cornac/models/diger/README.md @@ -0,0 +1,116 @@ +# DIGER + +Cornac implementation of **DIGER** (Differentiable Semantic ID for Generative +Recommendation, Fu et al., SIGIR 2026, +[arXiv:2601.19711](https://arxiv.org/abs/2601.19711)). A residual quantizer turns +item content embeddings into semantic IDs; a T5 generator predicts the next +item's ID. Frequency-gated Gumbel assignments let recommendation gradients +update the codebooks alongside the generator. + +## Requirements + +```bash +pip install -r cornac/models/diger/requirements.txt +``` + +## Usage + +Run the [Amazon Beauty example](../../../examples/diger_example.py) with +precomputed Sentence-T5 embeddings: + +```bash +python examples/diger_example.py --features beauty-st5.npz +``` + +The NPZ archive must contain `features` (a finite float32 matrix) and `ids` +(raw item IDs aligned with its rows), covering every item in the evaluation +splits. Add `--embedding openai` for precomputed OpenAI embeddings; no API key is +needed. Without `--features`, the example generates Sentence-T5 embeddings from +Cornac's item texts (requires `sentence-transformers`). These regenerated +features may differ from those used for the results below. + +Supply embeddings through `FeatureModality` in the Python API: + +```python +from cornac.data import FeatureModality +from cornac.eval_methods import NextItemEvaluation +from cornac.models import DIGER +from cornac.models.diger import BEAUTY_ST5_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 = DIGER(**{**BEAUTY_ST5_CONFIG, "device": "auto", "seed": 2025}) +``` + +`BEAUTY_ST5_CONFIG` and `BEAUTY_OPENAI_CONFIG` provide the selected Beauty +training recipe. Validation NDCG@10 is evaluated every epoch; training stops +after 15 non-improving epochs and restores the paired generator, tokenizer and +catalog. + +## Implementation notes + +- Training uses every history/next-item prefix; validation uses each session's + final target. +- A fresh fit pretrains a feature-width-matched RQ autoencoder (20 warmup epochs, + residual Lloyd initialization, 100 reconstruction/VQ epochs), then freezes its + encoder. This initialization adapts the released recipe. A bundle from + `training.pretrain_tokenizer` can be reused through `initializer` with the same + ordered features and tokenizer options. +- Joint training updates the codebooks and generator together; full-catalog + Sinkhorn assignments refresh each epoch. A deterministic suffix resolves + semantic-ID collisions. Unconstrained beam search scores valid full addresses; + set `n_beams` at least as large as the largest evaluation cutoff and no larger + than the codebook vocabulary. + +## Results + +The table reports CPU inference through the Cornac wrapper using six retained +checkpoints on Amazon Reviews 2014 Beauty: 12,068 items, 22,363 test queries, +beam 20 and history limit 210. Training was not rerun through this wrapper. +ST5 uses 768-dimensional Sentence-T5 features; OpenAI uses 3072-dimensional +`text-embedding-3-large` features. Values are mean ± sample SD across joint-training +seeds 2025/2026/2027 with a fixed seed-2024 initializer; metrics are fractions. + +| Features | Recall@5 | NDCG@5 | Recall@10 | NDCG@10 | Recall@20 | NDCG@20 | +| --- | --- | --- | --- | --- | --- | --- | +| ST5 | 0.041378 ± 0.001322 | 0.027401 ± 0.000342 | 0.063021 ± 0.001201 | 0.034355 ± 0.000252 | 0.093294 ± 0.003556 | 0.042015 ± 0.000854 | +| OpenAI | 0.044135 ± 0.001538 | 0.028961 ± 0.000630 | 0.069236 ± 0.001704 | 0.037026 ± 0.000674 | 0.105338 ± 0.001701 | 0.046100 ± 0.000367 | + +The retained features, initializers and checkpoints are not distributed with the +package; the presets alone do not reproduce this evaluation. + +### TIGER and RPG reference results + +For context, the following Beauty results are reported in the +[TIGER](../tiger/README.md#best-tiger-configs-vs-tuned-sasrec-beauty) and +[RPG](../rpg/README.md#graph-versus-exact-scoring) model documentation: + +| Model / recipe | Features | Recall@5 | NDCG@5 | Recall@10 | NDCG@10 | Recall@20 | NDCG@20 | +| --- | --- | --- | --- | --- | --- | --- | --- | +| TIGER / `PAISCHER_CONFIG`, RQVAE | ST5-base | 0.0419 | 0.0270 | 0.0634 | 0.0338 | 0.0953 | 0.0419 | +| TIGER / `PAISCHER_CONFIG`, rkmeans | OpenAI | 0.0405 | 0.0265 | 0.0678 | 0.0353 | 0.1045 | 0.0445 | +| RPG / `RPG_BEAUTY_CONFIG`, graph | OpenAI | 0.0533 | 0.0372 | 0.0789 | 0.0454 | 0.1113 | 0.0536 | + +These reference runs use seed 123; DIGER reports three-seed means above. +TIGER uses history length 20 and constrained beam 50; RPG uses history length +50 and graph beam 20, versus DIGER's history limit 210 and unconstrained beam +20. All use Amazon Reviews 2014 Beauty, but identical filtering, feature arrays +and evaluation populations have not been established. The rows provide context +rather than a controlled comparison. + +## Semantic ID evaluation + +Set `semantic_id_metrics=True` on `NextItemEvaluation`, or add +`--semantic-id-evaluation` to the example, to report SID and CCE Recall/NDCG +alongside ordinary item metrics from one inference. DIGER's complete semantic +IDs include a deterministic collision suffix, so each collision class contains +one item and all three metric views coincide in the reported runs. + +Invalid and repeated generated IDs are compacted for semantic-ID metrics. The +evaluator warns about short beams without expanding them. See the +[Semantic ID guide](../../../docs/source/api_ref/eval_methods.rst). diff --git a/cornac/models/diger/__init__.py b/cornac/models/diger/__init__.py new file mode 100644 index 00000000..7d796ce8 --- /dev/null +++ b/cornac/models/diger/__init__.py @@ -0,0 +1,21 @@ +# 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. +# ============================================================================ + +"""DIGER next-item recommendation.""" + +from .recom_diger import DIGER +from .diger_config import BEAUTY_ST5_CONFIG, BEAUTY_OPENAI_CONFIG + +__all__ = ["DIGER", "BEAUTY_ST5_CONFIG", "BEAUTY_OPENAI_CONFIG"] diff --git a/cornac/models/diger/diger.py b/cornac/models/diger/diger.py new file mode 100644 index 00000000..6cbc7092 --- /dev/null +++ b/cornac/models/diger/diger.py @@ -0,0 +1,311 @@ +# 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 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 details. + +"""DIGER computation; imported only when the optional backend is used. + +The public model module must not import this module at package import time. +This implementation uses PyTorch and the public Transformers T5 implementation; +the generator is initialized locally. +""" + +from collections import Counter +from dataclasses import dataclass + +import torch +from torch import nn +from torch.nn import functional as F +from transformers import T5Config, T5Model + + +@dataclass +class Quantization: + latent: torch.Tensor + quantized: torch.Tensor + probabilities: torch.Tensor + soft_probabilities: torch.Tensor + indices: torch.Tensor + vq_per_item: torch.Tensor + hot_masks: torch.Tensor + + +def hard_forward(soft, indices): + """Use integer assignments in forward, differentiable mixtures in backward.""" + hard = F.one_hot(indices, soft.shape[-1]).to(soft.dtype) + return hard - soft.detach() + soft + + +def balanced_assignments(distances, epsilon=0.003, iterations=50): + """Full-catalog balanced assignments on centered squared distances.""" + midpoint = (distances.max() + distances.min()) / 2 + amplitude = distances.max() - midpoint + 1e-5 + centered = ((distances - midpoint) / amplitude).double() + mass = torch.exp(-centered / epsilon) + mass = mass / mass.sum() + rows, columns = mass.shape + for _ in range(iterations): + mass = mass / mass.sum(dim=1, keepdim=True) / rows + mass = mass / mass.sum(dim=0, keepdim=True) / columns + if not torch.isfinite(mass).all(): + raise RuntimeError("Non-finite balanced code assignments") + return mass.argmax(dim=-1) + + +class SemanticTokenizer(nn.Module): + """Encoder-only residual quantizer with frequency-gated Gumbel assignments.""" + + def __init__(self, input_dim, hidden_dims=(2048, 1024, 512), latent_dim=256, + levels=3, codebook_size=256, beta=0.25, temperature=2.0, + usage_momentum=0.99, hot_ratio=1.5, noise_scale=1.0, + freeze_encoder=True): + super().__init__() + widths = (input_dim,) + tuple(hidden_dims) + (latent_dim,) + modules = [] + for index, (left, right) in enumerate(zip(widths[:-1], widths[1:])): + linear = nn.Linear(left, right) + nn.init.xavier_normal_(linear.weight) + nn.init.zeros_(linear.bias) + modules.append(linear) + if index < len(widths) - 2: + modules.append(nn.ReLU()) + self.encoder = nn.Sequential(*modules) + self.encoder.requires_grad_(not freeze_encoder) + self.codebooks = nn.ModuleList([ + nn.Embedding(codebook_size, latent_dim) for _ in range(levels) + ]) + for book in self.codebooks: + nn.init.uniform_(book.weight, -1 / codebook_size, 1 / codebook_size) + self.register_buffer("usage", torch.full((levels, codebook_size), 1 / codebook_size)) + self.beta = beta + self.temperature = temperature + self.usage_momentum = usage_momentum + self.hot_ratio = hot_ratio + self.noise_scale = noise_scale + + @staticmethod + def distances(residual, weight): + return (residual.square().sum(-1, keepdim=True) + + weight.square().sum(-1).unsqueeze(0) + - 2 * residual @ weight.t()) + + def forward(self, features, mode="frequency", noises=None): + if mode not in ("frequency", "plain", "fixed"): + raise ValueError("mode must be frequency, plain, or fixed") + latent = self.encoder(features) + residual = latent + vectors, mixtures, softs, indices, losses, masks = [], [], [], [], [], [] + for level, book in enumerate(self.codebooks): + logits = -self.distances(residual, book.weight) + nearest = logits.argmax(-1) + if mode == "frequency": + if self.training: + with torch.no_grad(): + frequency = torch.bincount(nearest, minlength=book.num_embeddings).to(self.usage.dtype) + frequency = frequency / frequency.sum().clamp_min(1) + self.usage[level].mul_(self.usage_momentum).add_(frequency, alpha=1-self.usage_momentum) + hot = self.usage[level][nearest] > self.hot_ratio / book.num_embeddings + else: + hot = torch.full_like(nearest, mode == "fixed", dtype=torch.bool) + hot = hot & self.training + noise = torch.zeros_like(logits) + # Consume one draw for every code even for cold items, preserving the + # assignment RNG contract while applying noise only to hot items. + if self.training and mode != "plain": + if noises is None: + uniform = torch.rand_like(logits) + eps = torch.finfo(uniform.dtype).eps + noise = -torch.log(-torch.log(uniform.clamp(eps, 1-eps))) + else: + noise = noises[level] + soft = F.softmax((logits + hot[:, None] * self.noise_scale * noise) / self.temperature, -1) + ids = soft.argmax(-1) + mixture = hard_forward(soft, ids) + vector = mixture @ book.weight + loss = ((vector-residual.detach()).square().mean(-1) + + self.beta * (vector.detach()-residual).square().mean(-1)) + vectors.append(vector) + mixtures.append(mixture) + softs.append(soft) + indices.append(ids) + losses.append(loss) + masks.append(hot) + residual = residual - vector + return Quantization(latent, torch.stack(vectors).sum(0), + torch.stack(mixtures, 1), torch.stack(softs, 1), + torch.stack(indices, 1), torch.stack(losses).mean(0), + torch.stack(masks, 1)) + + @torch.no_grad() + def catalog_codes(self, features, balanced=True, epsilon=0.003, iterations=50): + """Assign all catalog rows jointly; never run Sinkhorn on minibatches.""" + residual = self.encoder(features) + codes = [] + for book in self.codebooks: + distances = self.distances(residual, book.weight) + ids = (balanced_assignments(distances, epsilon, iterations) + if balanced else distances.argmin(-1)) + codes.append(ids) + residual = residual - book(ids) + return torch.stack(codes, 1) + + +def append_collision_suffix(prefixes, codebook_size=256): + """Injective full addresses in catalog row order, without renumbering rows.""" + counts = Counter() + suffix = [] + for prefix in prefixes.detach().cpu().tolist(): + key = tuple(prefix) + suffix.append(counts[key]) + counts[key] += 1 + if suffix and max(suffix) >= codebook_size: + raise ValueError("Collision bucket exceeds suffix vocabulary") + return torch.cat((prefixes, torch.tensor(suffix, device=prefixes.device, dtype=torch.long)[:, None]), 1) + + +class PositionalGenerator(nn.Module): + """T5 encoder/decoder with separate address-position embeddings and heads.""" + + def __init__(self, codebook_size=256, levels=3, model_dim=128, ff_dim=512, + num_layers=6, num_heads=4, key_dim=64, dropout=0.1): + super().__init__() + config = T5Config(vocab_size=1, d_model=model_dim, + d_ff=ff_dim, num_layers=num_layers, + num_decoder_layers=num_layers, num_heads=num_heads, + d_kv=key_dim, dropout_rate=dropout, + decoder_start_token_id=0, pad_token_id=0) + self.transformer = T5Model(config) + self.positions = nn.ModuleList([nn.Embedding(codebook_size, model_dim) + for _ in range(levels+1)]) + # DIGER initializes generator linear/embedding weights at std .02, + # including T5 components rather than retaining the generic T5 scale. + for module in self.modules(): + if isinstance(module, (nn.Linear, nn.Embedding)): + nn.init.normal_(module.weight, std=0.02) + if isinstance(module, nn.Linear) and module.bias is not None: + nn.init.zeros_(module.bias) + + def history_embeddings(self, probabilities, suffix, mask): + pieces = [probabilities[:, :, i] @ self.positions[i].weight + for i in range(len(self.positions)-1)] + pieces.append(self.positions[-1](suffix.clamp_min(0))) + embeddings = torch.stack(pieces, 2).flatten(1, 2) + expanded = mask[:, :, None].expand(-1, -1, len(self.positions)).flatten(1) + embeddings = torch.where(expanded[:, :, None], embeddings, + self.transformer.shared.weight[0].view(1, 1, -1)) + return embeddings, expanded + + def discrete_history_embeddings(self, codes, mask): + embeddings = torch.stack([embedding(codes[:, :, i].clamp_min(0)) + for i, embedding in enumerate(self.positions)], 2).flatten(1, 2) + expanded = mask[:, :, None].expand(-1, -1, len(self.positions)).flatten(1) + return torch.where(expanded[:, :, None], embeddings, + self.transformer.shared.weight[0].view(1, 1, -1)), expanded + + def decoder_embeddings(self, target_probabilities): + start = self.transformer.shared.weight[0].expand(target_probabilities.shape[0], -1) + return torch.stack([start] + [target_probabilities[:, i] @ self.positions[i].weight + for i in range(len(self.positions)-1)], 1) + + def forward(self, history_embeddings, mask, decoder_embeddings, encoder_outputs=None): + output = self.transformer(inputs_embeds=history_embeddings if encoder_outputs is None else None, + attention_mask=mask, decoder_inputs_embeds=decoder_embeddings, + encoder_outputs=encoder_outputs, output_hidden_states=True) + return torch.stack([output.last_hidden_state[:, i] @ self.positions[i].weight.t() + for i in range(decoder_embeddings.shape[1])], 1) + + @torch.no_grad() + def beam_search(self, codes, mask, beams=20, returns=None): + """Unconstrained address beams; scores are mean positional log probability.""" + returns = beams if returns is None else returns + if not 1 <= returns <= beams <= self.positions[0].num_embeddings: + raise ValueError("Require 1 <= returns <= beams <= vocabulary size") + batch = codes.shape[0] + embeddings, expanded = self.discrete_history_embeddings(codes, mask) + embeddings = embeddings.repeat_interleave(beams, 0) + expanded = expanded.repeat_interleave(beams, 0) + encoded = self.transformer.encoder(inputs_embeds=embeddings, attention_mask=expanded, return_dict=True) + paths = torch.zeros(batch*beams, 1, dtype=torch.long, device=codes.device) + scores = torch.full((batch, beams), -1e9, device=codes.device) + scores[:, 0] = 0 + vocabulary = self.positions[0].num_embeddings + for position in range(len(self.positions)): + decoder = torch.stack([self.transformer.shared(paths[:, 0])] + [ + self.positions[i](paths[:, i+1]) for i in range(position)], 1) + logits = self.forward(None, expanded, decoder, encoded)[:, -1] + candidates = (F.log_softmax(logits, -1) + scores.reshape(-1, 1)).reshape(batch, -1) + scores, choices = candidates.topk(beams, dim=1, sorted=True) + parents = choices // vocabulary + torch.arange(batch, device=codes.device)[:, None]*beams + paths = torch.cat([paths[parents.flatten()], (choices % vocabulary).reshape(-1, 1)], 1) + return (paths[:, 1:].reshape(batch, beams, -1)[:, :returns], + scores[:, :returns] / len(self.positions)) + + +class DIGERBackend(nn.Module): + """Joint computation with frozen input features and target-only auxiliary losses.""" + + def __init__(self, features, tokenizer_options=None, generator_options=None): + super().__init__() + if features.ndim != 2: + raise ValueError("features must have shape [catalog rows, semantic width]") + self.tokenizer_options = dict(tokenizer_options or {}) + self.generator_options = dict(generator_options or {}) + self.register_buffer("features", features.detach().clone()) + self.tokenizer = SemanticTokenizer(features.shape[1], **(tokenizer_options or {})) + options = dict(generator_options or {}) + options.setdefault("levels", len(self.tokenizer.codebooks)) + options.setdefault("codebook_size", self.tokenizer.codebooks[0].num_embeddings) + if (options["levels"] != len(self.tokenizer.codebooks) + or options["codebook_size"] != self.tokenizer.codebooks[0].num_embeddings): + raise ValueError("Generator levels and codebook_size must match tokenizer") + self.generator = PositionalGenerator(**options) + + def initialize_tokenizer(self, state, semantic_width): + """Strictly load pretrained encoder/books/frequency state. + + A tokenizer trained for a different feature width cannot initialize this + model. The fit controller must require a matched initializer or train + one before freezing the encoder; random encoder construction is useful + only for component tests. + """ + if semantic_width != self.features.shape[1]: + raise ValueError("Tokenizer semantic feature width differs from catalog") + self.tokenizer.load_state_dict(state, strict=True) + + def batch_loss(self, history_ids, target_ids, catalog, mode="frequency", + cached_forward=True, code_weight=1.0, reconstruction_weight=1.0, + vq_weight=1.0, noises=None, detach_recommendation=False): + """Catalog row0 is padding; catalog codes at padding are all -1.""" + mask = history_ids != 0 + targets = target_ids.reshape(-1) + unique, inverse = torch.unique(torch.cat([history_ids[mask], targets]), sorted=True, return_inverse=True) + quantized = self.tokenizer(self.features[unique], mode=mode, noises=noises) + probabilities = (hard_forward(quantized.soft_probabilities, catalog[unique, :-1]) + if cached_forward else quantized.probabilities) + if detach_recommendation: + probabilities = probabilities.detach() + history_inverse = torch.zeros_like(history_ids) + history_inverse[mask] = inverse[:mask.sum()] + target_inverse = inverse[mask.sum():] + history_probabilities = probabilities[history_inverse] + target_probabilities = probabilities[target_inverse] + embeddings, expanded = self.generator.history_embeddings(history_probabilities, catalog[history_ids, -1], mask) + logits = self.generator(embeddings, expanded, self.generator.decoder_embeddings(target_probabilities)) + denominator = targets.shape[0]*len(self.generator.positions) + semantic = -(target_probabilities * F.log_softmax(logits[:, :-1].float(), -1)).sum()/denominator + suffix = F.cross_entropy(logits[:, -1].float(), catalog[targets, -1], reduction="sum")/denominator + reconstruction = (quantized.quantized-quantized.latent).square().mean(-1)[target_inverse].mean() + vq = quantized.vq_per_item[target_inverse].mean() + loss = code_weight*(semantic+suffix) + reconstruction_weight*reconstruction + vq_weight*vq + return {"loss": loss, "semantic": semantic, "suffix": suffix, + "reconstruction": reconstruction, "vq": vq, "logits": logits, + "quantization": quantized, "unique_ids": unique} + + @torch.no_grad() + def catalog(self, balanced=True, epsilon=0.003, iterations=50): + prefixes = self.tokenizer.catalog_codes(self.features[1:], balanced, epsilon, iterations) + full = append_collision_suffix(prefixes, self.generator.positions[-1].num_embeddings) + return torch.cat([full.new_full((1, full.shape[1]), -1), full], 0) diff --git a/cornac/models/diger/diger_config.py b/cornac/models/diger/diger_config.py new file mode 100644 index 00000000..ec96f324 --- /dev/null +++ b/cornac/models/diger/diger_config.py @@ -0,0 +1,40 @@ +# 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. +# ============================================================================ + +"""Beauty validation-selected recipes for retained ST5/OpenAI item embeddings. + +Both conditions selected the same learning rates from six candidates. The +retained study used a seed-2024 initializer and joint seeds 2025/2026/2027. +These dictionaries train a fresh initializer unless one is explicitly supplied; +they do not reproduce the retained study assets by themselves. See README.md. +""" + +BEAUTY_ST5_CONFIG = { + "n_epochs": 120, + "batch_size": 256, + "eval_batch_size": 32, + "learning_rate": 5e-4, + "tokenizer_learning_rate": 3e-5, + "weight_decay": 0.05, + "warmup_steps": 500, + "max_len": 210, + "n_beams": 20, + "early_stopping_patience": 15, + "initializer_seed": 2024, + "seed": 2025, +} + +# Architecture and selected rates coincide; feature widths differ (768/3072). +BEAUTY_OPENAI_CONFIG = dict(BEAUTY_ST5_CONFIG) diff --git a/cornac/models/diger/recom_diger.py b/cornac/models/diger/recom_diger.py new file mode 100644 index 00000000..4a421e6f --- /dev/null +++ b/cornac/models/diger/recom_diger.py @@ -0,0 +1,249 @@ +# 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. +# ============================================================================ + +"""DIGER next-item recommendation with T5 and residual quantization.""" + +from pathlib import Path + +import numpy as np + +from ..recommender import NextItemRecommender, SemanticIDMixin + + +class DIGER(NextItemRecommender, SemanticIDMixin): + """DIGER with frequency-gated residual quantization and T5 generation. + + Supply precomputed item embeddings through ``FeatureModality`` in + ``NextItemEvaluation``. Features are used without additional normalization. + Training expands each session into all next-item prefixes; validation uses + only each session's final target. Sessions shorter than two items are skipped. + + Parameters + ---------- + name: str, default: "DIGER" + Model name. + initializer: dict or path, optional + Bundle produced by ``training.pretrain_tokenizer`` on the same ordered + catalog/features. If absent, train a matched-feature RQ autoencoder first. + tokenizer_options: dict, optional + Options for ``diger.SemanticTokenizer``, including hidden_dims, latent_dim, + levels and codebook_size. The encoder is frozen after initialization; + recommendation gradients update the residual codebooks. + generator_options: dict, optional + Options for ``diger.PositionalGenerator``, including model_dim, ff_dim, + num_layers and num_heads. Levels and vocabulary must match the tokenizer. + initializer_options: dict, optional + ``pretrain_tokenizer`` options (default: 100 epochs, 20 warmup epochs, + batch size 256, learning rate 0.001). + initializer_seed: int, default: 2024 + Random seed for pretraining a fresh initializer. + n_epochs: int, default: 120 + Joint training horizon; validation NDCG@10 selects paired model/catalog. + batch_size: int, default: 256 + Number of training prefixes per batch. + eval_batch_size: int, default: 32 + Number of validation sessions per batch. + learning_rate: float, default: 0.0005 + Generator AdamW learning rate. + tokenizer_learning_rate: float, default: 0.00003 + Codebook AdamW learning rate. Both rates were selected on Beauty validation. + weight_decay: float, default: 0.05 + Weight decay for both joint-training AdamW optimizers. + warmup_steps: int, default: 500 + Linear warmup steps before cosine learning-rate decay. + max_len: int, default: 210 + Maximum number of history items, used consistently during fit/inference. + n_beams: int, default: 20 + Unconstrained positional beam width. Use at least the evaluation cutoff. + early_stopping_patience: int, default: 15 + Stop after this many epochs without strict validation improvement. + device: str, default: 'cpu' + PyTorch device, or 'auto' to select CUDA when available. + trainable: bool, default: True + Whether ``fit`` trains the model. + verbose: bool, default: False + Whether to print epoch loss and validation scores. + seed: int, default: 2024 + Random seed for generator initialization and joint training. + """ + + def __init__(self, name="DIGER", initializer=None, tokenizer_options=None, + generator_options=None, initializer_options=None, + initializer_seed=2024, n_epochs=120, batch_size=256, + eval_batch_size=32, learning_rate=5e-4, + tokenizer_learning_rate=3e-5, weight_decay=0.05, + warmup_steps=500, max_len=210, n_beams=20, + early_stopping_patience=15, + device="cpu", trainable=True, verbose=False, seed=2024): + super().__init__(name=name, trainable=trainable, verbose=verbose) + self.initializer = initializer + self.tokenizer_options = dict(tokenizer_options or {}) + self.generator_options = dict(generator_options or {}) + if min(n_epochs, batch_size, eval_batch_size, n_beams) < 1 or max_len < 2: + raise ValueError("Epochs, batch sizes, max_len and n_beams must be positive") + if min(learning_rate, tokenizer_learning_rate) <= 0 or weight_decay < 0: + raise ValueError("Learning rates must be positive and weight_decay nonnegative") + if warmup_steps < 0 or early_stopping_patience < 1: + raise ValueError("warmup_steps must be nonnegative and patience positive") + self.initializer_options = dict(initializer_options or {}) + self.initializer_seed = initializer_seed + self.n_epochs = n_epochs + self.batch_size = batch_size + self.eval_batch_size = eval_batch_size + self.learning_rate = learning_rate + self.tokenizer_learning_rate = tokenizer_learning_rate + self.weight_decay = weight_decay + self.warmup_steps = warmup_steps + self.max_len = max_len + self.n_beams = n_beams + self.early_stopping_patience = early_stopping_patience + self.seed = seed + self.device = device + + def _rows(self, dataset): + rows = [] + excluded = getattr(dataset, "catalog_only_session_indices", set()) + for session, indices in dataset.sessions.items(): + if session in excluded: + continue + items = [int(i) + 1 for i in dataset.uir_tuple[1][indices]] + stops = range(1, len(items)) if dataset is self.train_set else [len(items)-1] + for stop in stops: + if stop > 0: + rows.append((items[max(0, stop-self.max_len):stop], items[stop])) + if not rows: + raise ValueError("DIGER requires sessions containing history and target") + return rows + + def _feature_rows(self, train_set): + modality = getattr(train_set, "item_feature", None) + features = getattr(modality, "features", None) + if features is None: + raise ValueError("DIGER requires item_feature=FeatureModality(features=..., ids=...)") + features = np.asarray(features, dtype="float32") + ordered_ids = [raw for raw, _ in sorted(train_set.iid_map.items(), key=lambda pair: pair[1])] + if (features.ndim != 2 or features.shape[0] < len(ordered_ids) + or features.shape[1] == 0 or not np.isfinite(features).all()): + raise ValueError("Require finite 2D features covering every Cornac item") + ids = getattr(modality, "ids", None) + if ids is not None and list(ids[:len(ordered_ids)]) != ordered_ids: + raise ValueError("FeatureModality IDs must cover and align with every Cornac item") + return features[:len(ordered_ids)].copy(), ordered_ids + + def fit(self, train_set, val_set=None): + super().fit(train_set, val_set) + if not self.trainable: + return self + if val_set is None: + raise ValueError("DIGER requires validation sessions for model selection") + features, _ = self._feature_rows(train_set) + train_rows, val_rows = self._rows(train_set), self._rows(val_set) + import torch + from .diger import DIGERBackend + from .training import pretrain_tokenizer, tensor_hash, train + + self._ensure_device(torch) + padded = np.vstack([np.zeros((1, features.shape[1]), dtype="float32"), features]) + features = torch.tensor(padded, device=self.device) + initializer = self.initializer + if initializer is None: + initializer = pretrain_tokenizer( + features, self.tokenizer_options, + **{"seed": self.initializer_seed, **self.initializer_options}) + elif isinstance(initializer, (str, Path)): + initializer = torch.load(initializer, map_location="cpu", weights_only=False) + if initializer.get("format") != "cornac-diger-initializer-v1": + raise ValueError("Require a matched pretrained tokenizer initializer") + expected = {k: v for k, v in initializer["tokenizer_options"].items() + if k != "freeze_encoder"} + actual = {k: v for k, v in self.tokenizer_options.items() if k != "freeze_encoder"} + if expected != actual: + raise ValueError("Initializer tokenizer recipe differs from tokenizer options") + if initializer["feature_hash"] != tensor_hash(features): + raise ValueError("Initializer feature hash differs from ordered catalog") + torch.manual_seed(self.seed) + self.backend_ = DIGERBackend(features, self.tokenizer_options, + self.generator_options).to(self.device) + if self.n_beams > self.backend_.generator.positions[0].num_embeddings: + raise ValueError("n_beams must not exceed codebook_size") + self.backend_.initialize_tokenizer(initializer["state"], initializer["semantic_width"]) + config = {"epochs": self.n_epochs, "batch_size": self.batch_size, + "eval_batch_size": self.eval_batch_size, "lr_rec": self.learning_rate, + "lr_id": self.tokenizer_learning_rate, "weight_decay": self.weight_decay, + "warmup_steps": self.warmup_steps, "max_length": self.max_len, + "beams": self.n_beams, "patience": self.early_stopping_patience, + "seed": self.seed} + selected, self.training_summary_ = train( + self.backend_, train_rows, val_rows, config=config, verbose=self.verbose) + self.best_epoch, self.best_value = selected["epoch"], selected["ndcg@10"] + self.codes_ = selected["catalog"].numpy().copy() + self.sid_to_item_ = {tuple(code): item for item, code in enumerate(self.codes_[1:])} + self.backend_.to("cpu") + return self + + def _ensure_device(self, torch): + if self.device == "auto": + self.device = "cuda" if torch.cuda.is_available() else "cpu" + elif torch.device(self.device).type == "cuda" and not torch.cuda.is_available(): + self.device = "cpu" + + def save(self, save_dir=None, save_trainset=False, metadata=None): + if save_dir is None or not hasattr(self, "backend_"): + return super().save(save_dir, save_trainset, metadata) + original_device = self.backend_.features.device + self.backend_.to("cpu") + try: + return super().save(save_dir, save_trainset, metadata) + finally: + self.backend_.to(original_device) + + def semantic_id_catalog(self): + """Return complete collision-suffixed IDs in Cornac item order.""" + if not hasattr(self, "codes_"): + raise ValueError("DIGER must be fitted before exporting its SID catalog") + return self.codes_[1:] + + def _beam_predictions(self, history_items): + """Decode raw addresses once, retaining invalid and duplicate rows.""" + import torch + + catalog = self.semantic_id_catalog() + history = list(history_items)[-self.max_len:] + if not history: + return np.empty((0, catalog.shape[1]), dtype="int64"), np.empty(0) + if any(i < 0 or i >= self.total_items for i in history): + raise ValueError("History item outside catalog") + self._ensure_device(torch) + self.backend_.to(self.device).eval() + codes = torch.tensor(self.codes_[np.asarray(history, dtype="int64") + 1], + device=self.device)[None] + paths, values = self.backend_.generator.beam_search( + codes, torch.ones(codes.shape[:2], dtype=torch.bool, device=self.device), + beams=self.n_beams) + return paths[0].cpu().numpy(), values[0].cpu().numpy() + + def _score_with_semantic_ids(self, user_idx, history_items, **kwargs): + addresses, values = self._beam_predictions(history_items) + scores = np.full(self.total_items, -np.inf, dtype="float64") + for address, value in zip(addresses, values): + item = self.sid_to_item_.get(tuple(address)) + if item is not None: + scores[item] = value + return scores, addresses + + def score(self, user_idx, history_items, item_idx=None, **kwargs): + """Return catalog scores from unconstrained full-address beam search.""" + scores, _ = self._score_with_semantic_ids(user_idx, history_items, **kwargs) + return scores if item_idx is None else scores[item_idx] diff --git a/cornac/models/diger/requirements.txt b/cornac/models/diger/requirements.txt new file mode 100644 index 00000000..b3a896ba --- /dev/null +++ b/cornac/models/diger/requirements.txt @@ -0,0 +1,2 @@ +torch>=2.5.1 +transformers>=4.57.1 diff --git a/cornac/models/diger/training.py b/cornac/models/diger/training.py new file mode 100644 index 00000000..1616d226 --- /dev/null +++ b/cornac/models/diger/training.py @@ -0,0 +1,201 @@ +# 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. +# ============================================================================ + +"""Matched-feature tokenizer initialization and joint DIGER training.""" +import hashlib +import math +import random + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F +from .diger import SemanticTokenizer + + +def tensor_hash(value): + array = value.detach().cpu().contiguous().numpy() + return hashlib.sha256(str(array.shape).encode() + array.tobytes()).hexdigest() + + +def pretrain_tokenizer(features, options=None, *, epochs=100, batch_size=256, + learning_rate=1e-3, seed=2024, + warmup_epochs=20): + """Learn a matched-width RQ autoencoder initializer on catalog features. + + This explicitly declared shared-feature initializer uses deterministic nearest + residual codes, straight-through reconstruction, and codebook/commitment MSE. + It is an adaptation, not a claim to reproduce an external released initializer. + Padding is excluded. The decoder is discarded before recommendation training. + """ + if features.ndim != 2 or len(features) < 2 or not torch.isfinite(features).all(): + raise ValueError("Require finite features with padding plus catalog rows") + if epochs < 1 or batch_size < 1: + raise ValueError("epochs and batch_size must be positive") + torch.manual_seed(seed) + options = dict(options or {}) + options['freeze_encoder'] = False + tokenizer = SemanticTokenizer(features.shape[1], **options).to(features.device) + widths = (options.get('latent_dim', 256),) + tuple(reversed(options.get('hidden_dims', (2048, 1024, 512)))) + (features.shape[1],) + layers = [] + for index, (left, right) in enumerate(zip(widths[:-1], widths[1:])): + layers.append(nn.Linear(left, right)) + if index < len(widths)-2: + layers.append(nn.ReLU()) + decoder = nn.Sequential(*layers).to(features.device) + optimizer = torch.optim.AdamW(list(tokenizer.parameters()) + list(decoder.parameters()), lr=learning_rate) + # Reconstruction warmup learns a nonconstant feature representation before + # nearest-code training. Residual Lloyd initialization avoids trapping all + # normalized semantic vectors in the same tiny random code vector. + warmup_optimizer = torch.optim.AdamW(list(tokenizer.encoder.parameters()) + list(decoder.parameters()), lr=learning_rate) + reconstruction_scale = features[1:].var(dim=0, unbiased=False).mean().clamp_min(1e-8) + warmup_losses = [] + for epoch in range(warmup_epochs): + total = 0. + for ids in (torch.randperm(len(features)-1, device=features.device)+1).split(batch_size): + loss = F.mse_loss(decoder(tokenizer.encoder(features[ids])), features[ids])/reconstruction_scale + warmup_optimizer.zero_grad(); loss.backward(); warmup_optimizer.step() + total += float(loss.detach())*len(ids) + warmup_losses.append(total/(len(features)-1)) + with torch.no_grad(): + residual = tokenizer.encoder(features[1:]) + for book in tokenizer.codebooks: + centers = residual[torch.randint(len(residual), (book.num_embeddings,), device=features.device)].clone() + for _ in range(20): + assignments = tokenizer.distances(residual, centers).argmin(-1) + counts = torch.bincount(assignments, minlength=len(centers)) + sums = torch.zeros_like(centers).index_add_(0, assignments, residual) + occupied = counts > 0 + centers[occupied] = sums[occupied]/counts[occupied, None] + book.weight.copy_(centers) + residual = residual-centers[assignments] + losses = [] + for epoch in range(epochs): + order = torch.randperm(len(features)-1, device=features.device)+1 + total = 0. + for ids in order.split(batch_size): + quantized = tokenizer(features[ids], mode='plain') + straight = quantized.latent + (quantized.quantized-quantized.latent).detach() + loss = F.mse_loss(decoder(straight), features[ids])/reconstruction_scale + quantized.vq_per_item.mean() + optimizer.zero_grad() + loss.backward() + optimizer.step() + total += float(loss.detach()) * len(ids) + losses.append(total/(len(features)-1)) + tokenizer.eval(); decoder.eval() + with torch.no_grad(): + final_reconstruction = 0. + for values in features[1:].split(batch_size): + quantized = tokenizer(values, mode='plain') + final_reconstruction += float(F.mse_loss(decoder(quantized.quantized), values)/reconstruction_scale)*len(values) + final_reconstruction /= len(features)-1 + bundle = {'format': 'cornac-diger-initializer-v1', 'semantic_width': features.shape[1], + 'feature_hash': tensor_hash(features), 'tokenizer_options': options, + 'terminal_normalized_reconstruction': final_reconstruction, + 'mean_feature_baseline_normalized_reconstruction': 1.0, + 'pretrain_config': {'epochs': epochs, 'warmup_epochs': warmup_epochs, 'batch_size': batch_size, 'learning_rate': learning_rate, 'weight_decay': .01, 'lloyd_iterations': 20}, + 'state': {k: v.detach().cpu() for k, v in tokenizer.state_dict().items()}, + 'losses': losses, 'warmup_losses': warmup_losses, 'reconstruction_variance_scale': float(reconstruction_scale), 'warmup_epochs': warmup_epochs, 'seed': seed, + 'recipe': 'matched-feature AE warmup + residual Lloyd20 + variance-scaled nearest-code ST reconstruction + VQ'} + return bundle + + +def batch(rows, indices, device, max_length): + selected = [rows[int(i)] for i in indices] + histories = [history[-max_length:] + for history, _ in selected] + padded = torch.zeros(len(selected), max(map(len, histories)), dtype=torch.long, device=device) + for index, history in enumerate(histories): + padded[index, :len(history)] = torch.tensor(history, device=device) + return padded, torch.tensor([target for _, target in selected], device=device) + + +def validate_rows(rows, catalog_size): + if not rows: + raise ValueError('Training and validation examples must be nonempty') + for history, target in rows: + if not history or any(not 1 <= int(i) <= catalog_size for i in history) or not 1 <= int(target) <= catalog_size: + raise ValueError('Example contains empty history, padding, or unknown item') + + +def train(backend, train_rows, validation_rows, *, config, verbose=False): + """Select generator, tokenizer and catalog together by validation NDCG@10.""" + validate_rows(train_rows, len(backend.features)-1) + validate_rows(validation_rows, len(backend.features)-1) + device = backend.features.device + steps = math.ceil(len(train_rows)/config['batch_size'])*config['epochs'] + parameters = list(backend.tokenizer.codebooks.parameters()) + backend.tokenizer.requires_grad_(True) + backend.tokenizer.encoder.requires_grad_(False) + rec = torch.optim.AdamW(backend.generator.parameters(), lr=config['lr_rec'], weight_decay=config['weight_decay']) + ids = torch.optim.AdamW(parameters, lr=config['lr_id'], weight_decay=config['weight_decay']) + def schedule(step): + warmup = config['warmup_steps'] + if step < warmup: + return step/max(1, warmup) + return max(0., .5*(1+math.cos(math.pi*(step-warmup)/max(1, steps-warmup)))) + schedulers = [torch.optim.lr_scheduler.LambdaLR(opt, schedule) for opt in (rec, ids)] + global_step, best_score, best_epoch, stale = 0, -float('inf'), None, 0 + best_state = None + history = [] + random.seed(config['seed']); np.random.seed(config['seed']); torch.manual_seed(config['seed']) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(config['seed']) + catalog = backend.catalog() + for epoch in range(config['epochs']): + backend.train() + total = 0. + order = torch.randperm(len(train_rows)) + for indices in order.split(config['batch_size']): + h, t = batch(train_rows, indices, device, config['max_length']) + rec.zero_grad(); ids.zero_grad() + result = backend.batch_loss(h, t, catalog) + result['loss'].backward() + nn.utils.clip_grad_norm_(backend.generator.parameters(), 1.) + nn.utils.clip_grad_norm_(parameters, 1.) + rec.step() + ids.step() + for scheduler in schedulers: + scheduler.step() + global_step += 1 + total += float(result['loss'].detach())*len(indices) + backend.eval() + catalog = backend.catalog() + ndcg = 0. + with torch.no_grad(): + for indices in torch.arange(len(validation_rows)).split(config['eval_batch_size']): + h, t = batch(validation_rows, indices, device, config['max_length']) + paths, _ = backend.generator.beam_search(catalog[h], h != 0, beams=config['beams']) + matches = (paths[:, :10] == catalog[t, None]).all(-1) + discounts = 1/torch.log2(torch.arange(matches.shape[1], device=device)+2.) + ndcg += float((matches*discounts).sum()) + score = ndcg/len(validation_rows) + history.append({'epoch': epoch, 'loss': total/len(train_rows), 'ndcg@10': score}) + if best_epoch is None or score > best_score: + best_score, best_epoch, stale = score, epoch, 0 + best_state = {'state': {k: v.detach().cpu().clone() for k, v in backend.state_dict().items()}, + 'catalog': catalog.detach().cpu(), 'epoch': epoch, 'ndcg@10': score} + else: + stale += 1 + if verbose: + print(f'DIGER epoch {epoch}: loss={history[-1]["loss"]:.6f}, ndcg@10={score:.6f}', flush=True) + if stale >= config['patience']: + break + if best_state is None: + raise ValueError('No training epochs remain') + backend.load_state_dict(best_state['state'], strict=True) + backend.eval() + return best_state, {'best_epoch': best_epoch, 'best_score': best_score, + 'global_step': global_step, 'history': history} diff --git a/docs/source/api_ref/models.rst b/docs/source/api_ref/models.rst index fddac1a1..cab9c242 100644 --- a/docs/source/api_ref/models.rst +++ b/docs/source/api_ref/models.rst @@ -338,3 +338,8 @@ Semantic ID Model Contracts -------------------------- .. autoclass:: cornac.models.recommender.SemanticIDMixin :members: + +DIGER +----- +.. automodule:: cornac.models.diger.recom_diger + :members: diff --git a/examples/README.md b/examples/README.md index 44a1be0b..2db44bc2 100644 --- a/examples/README.md +++ b/examples/README.md @@ -136,6 +136,8 @@ [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. +[diger_example.py](diger_example.py) - Train DIGER with T5 generation and frequency-gated residual quantization on Amazon Beauty through FeatureModality, using validation-selected ST5/OpenAI presets. + [diffgrm_example.py](diffgrm_example.py) - Masked-diffusion generation of PSE semantic IDs (DiffGRM) on Amazon Sports with Sentence-T5 item content embeddings, OCN training, and CPD decoding. ---- diff --git a/examples/diger_example.py b/examples/diger_example.py new file mode 100644 index 00000000..28990e6e --- /dev/null +++ b/examples/diger_example.py @@ -0,0 +1,81 @@ +"""DIGER on Amazon Beauty through Cornac's normal next-item API. + +Run ``python examples/diger_example.py --features beauty-st5.npz`` with an +archive containing ``features`` (finite float32 matrix) and ``ids`` (raw item IDs). +Add ``--embedding openai`` for precomputed OpenAI embeddings; no API key is +needed. Without an archive, generate Sentence-T5 features from Cornac's item +texts (requires sentence-transformers). Results depend on the supplied features; +see the model README for the evaluated setup. +""" + +import argparse + +import numpy as np + +import cornac +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 DIGER +from cornac.models.diger import BEAUTY_OPENAI_CONFIG, BEAUTY_ST5_CONFIG + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--features", help="NPZ containing features and raw ids") + parser.add_argument("--embedding", choices=("st5", "openai"), default="st5") + parser.add_argument("--initializer", help="Optional matched initializer bundle") + parser.add_argument("--device", default="auto") + parser.add_argument( + "--semantic-id-evaluation", + action="store_true", + help="Report opt-in SID and collision-corrected metrics", + ) + args = parser.parse_args() + if args.embedding == "openai" and not args.features: + parser.error("OpenAI mode requires precomputed --features") + + data = amazon_review.load_feedback(category="beauty") + if args.features: + with np.load(args.features, allow_pickle=False) as archive: + features, item_ids = archive["features"], archive["ids"].tolist() + else: + from sentence_transformers import SentenceTransformer + + texts, item_ids = amazon_review.load_text(category="beauty") + encoder = SentenceTransformer("sentence-t5-base") + features = encoder.encode(texts, batch_size=64, show_progress_bar=True) + del encoder + + evaluation = NextItemEvaluation.leave_last_out( + data=data, + exclude_unknowns=True, + verbose=True, + semantic_id_metrics=args.semantic_id_evaluation, + item_feature=FeatureModality(features=features, ids=item_ids), + ) + preset = BEAUTY_ST5_CONFIG if args.embedding == "st5" else BEAUTY_OPENAI_CONFIG + model = DIGER( + **preset, + initializer=args.initializer, + device=args.device, + verbose=True, + ) + metrics = [ + Recall(k=5), + Recall(k=10), + Recall(k=20), + NDCG(k=5), + NDCG(k=10), + NDCG(k=20), + ] + cornac.Experiment( + eval_method=evaluation, + models=[model], + metrics=metrics, + ).run() + + +if __name__ == "__main__": + main() diff --git a/tests/cornac/models/test_diger.py b/tests/cornac/models/test_diger.py new file mode 100644 index 00000000..21f1488d --- /dev/null +++ b/tests/cornac/models/test_diger.py @@ -0,0 +1,159 @@ +# Copyright 2026 The Cornac Authors. All Rights Reserved. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may obtain a copy at http://www.apache.org/licenses/LICENSE-2.0 + +"""Small CPU checks for DIGER training and its optional semantic-ID adapter.""" + +import os +import subprocess +import warnings +import sys +from unittest.mock import patch + +import numpy as np +import pytest + +from cornac.data import FeatureModality +from cornac.eval_methods import NextItemEvaluation +from cornac.metrics import Recall +from cornac.models import DIGER + + +def tiny_evaluation(semantic_id_metrics=False): + return NextItemEvaluation.from_splits( + train_data=[("u", "tr", "a", 0), ("u", "tr", "b", 1), + ("u", "tr", "c", 2)], + val_data=[("u", "va", "b", 3), ("u", "va", "a", 4), + ("u", "va", "c", 5)], + test_data=[("u", "te", "a", 6), ("u", "te", "b", 7)], + fmt="USIT", mode="last", semantic_id_metrics=semantic_id_metrics, + item_feature=FeatureModality(features=np.eye(3), ids=["a", "b", "c"]), + ) + + +def tiny_model(**kwargs): + return DIGER( + tokenizer_options={"hidden_dims": (4,), "latent_dim": 2, + "levels": 1, "codebook_size": 4}, + generator_options={"model_dim": 4, "ff_dim": 8, "num_layers": 1, + "num_heads": 1, "key_dim": 4, "dropout": 0}, + initializer_options={"epochs": 1, "warmup_epochs": 1, "batch_size": 3}, + n_epochs=1, batch_size=2, eval_batch_size=1, n_beams=3, + max_len=2, warmup_steps=0, seed=7, initializer_seed=7, **kwargs, + ) + + +@pytest.fixture(scope="module") +def fitted(): + torch = pytest.importorskip("torch") + pytest.importorskip("transformers") + old_threads = torch.get_num_threads() + torch.set_num_threads(1) + evaluation = tiny_evaluation() + model = tiny_model().fit(evaluation.train_set, evaluation.val_set) + yield model, evaluation + torch.set_num_threads(old_threads) + + +def test_optional_dependencies_load_only_when_used(): + script = """ +import sys +class BlockBackend: + def find_spec(self, fullname, *args): + if fullname.split('.')[0] in ('torch', 'transformers'): + raise ImportError('optional backend blocked') +sys.meta_path.insert(0, BlockBackend()) +from cornac.models import DIGER +DIGER() +assert 'torch' not in sys.modules +assert 'transformers' not in sys.modules +""" + subprocess.run([sys.executable, "-c", script], check=True, env=os.environ.copy()) + + +def test_prefixes_and_feature_alignment(): + evaluation = tiny_evaluation() + model = DIGER(trainable=False).fit(evaluation.train_set, evaluation.val_set) + assert model._rows(evaluation.train_set) == [([1], 2), ([1, 2], 3)] + assert model._rows(evaluation.val_set) == [([2, 1], 3)] + features, ids = model._feature_rows(evaluation.train_set) + np.testing.assert_array_equal(features, np.eye(3)) + assert ids == ["a", "b", "c"] + evaluation.train_set.item_feature.ids = ["c", "b", "a"] + with pytest.raises(ValueError, match="align"): + model._feature_rows(evaluation.train_set) + + +def test_catalog_and_paired_checkpoint(fitted): + model, _ = fitted + np.testing.assert_array_equal(model.semantic_id_catalog(), model.codes_[1:]) + assert model.semantic_id_catalog().shape == (3, 2) + assert len(set(map(tuple, model.semantic_id_catalog()))) == 3 + np.testing.assert_array_equal(model.backend_.catalog().cpu(), model.codes_) + assert model.best_epoch == model.training_summary_["best_epoch"] + assert all(np.isfinite(row["loss"]) for row in model.training_summary_["history"]) + + +def test_shared_inference_keeps_raw_beam_rows(fitted): + import torch + + model, _ = fitted + valid = model.codes_[1].copy() + raw = np.stack([valid, [99, 99], valid]) + paths = torch.tensor(raw)[None] + values = torch.tensor([[-1., -2., -3.]]) + with patch.object(model.backend_.generator, "beam_search", return_value=(paths, values)) as beam: + ranks, scores, semantic_ids = model.rank_with_semantic_ids(0, history_items=[0], k=2) + assert beam.call_count == 1 + assert beam.call_args.kwargs["beams"] == model.n_beams + np.testing.assert_array_equal(semantic_ids, raw) + assert scores[0] == -3 # Preserve the native address-to-item score mapping. + assert np.isneginf(scores[1:]).all() + assert ranks[0] == 0 + + +def test_empty_and_unknown_histories(fitted): + model, _ = fitted + scores, ids = model._score_with_semantic_ids(0, []) + assert np.isneginf(scores).all() + assert ids.shape == (0, 2) + with pytest.raises(ValueError, match="outside catalog"): + model.score(0, [model.total_items]) + + +def test_save_load_and_cpu_fallback(fitted, tmp_path): + import torch + + model, _ = fitted + before = model.score(0, [0, 1]) + path = model.save(str(tmp_path)) + loaded = DIGER.load(path) + loaded.device = "cuda" + with patch.object(torch.cuda, "is_available", return_value=False): + np.testing.assert_array_equal(loaded.score(0, [0, 1]), before) + assert loaded.device == "cpu" + np.testing.assert_array_equal(loaded.codes_, model.codes_) + + +def test_fit_reproducibility_and_opt_in_evaluation(fitted): + model, _ = fitted + evaluation = tiny_evaluation(semantic_id_metrics=True) + repeated = tiny_model().fit(evaluation.train_set, evaluation.val_set) + np.testing.assert_array_equal(repeated.codes_, model.codes_) + np.testing.assert_array_equal(repeated.score(0, [0]), model.score(0, [0])) + repeated.trainable = False + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + result, _ = evaluation.evaluate(repeated, [Recall(k=1)], user_based=False) + assert all("already fitted" in str(w.message) or + "valid distinct Semantic IDs" in str(w.message) for w in caught) + assert {"Recall@1", "SID-Recall@1", "CCE-Recall@1"} <= set(result.metric_avg_results) + + +def test_decoder_rejects_oversized_beams(fitted): + import torch + + model, _ = fitted + codes = torch.tensor(model.codes_[[1]])[None] + with pytest.raises(ValueError, match="vocabulary"): + model.backend_.generator.beam_search(codes, torch.ones((1, 1), dtype=torch.bool), beams=5) From f121b111ab7e7c8d4769512f6a3591500727f23e Mon Sep 17 00:00:00 2001 From: hieuddo Date: Mon, 2 Nov 2026 01:26:38 +0800 Subject: [PATCH 2/2] docs: compare DIGER with TIGER and RPG --- cornac/models/diger/README.md | 38 +++++++++++++++++------------------ 1 file changed, 19 insertions(+), 19 deletions(-) diff --git a/cornac/models/diger/README.md b/cornac/models/diger/README.md index 3d35c7e4..c77605a2 100644 --- a/cornac/models/diger/README.md +++ b/cornac/models/diger/README.md @@ -69,39 +69,39 @@ catalog. ## Results -The table reports CPU inference through the Cornac wrapper using six retained +The DIGER rows report CPU inference through the Cornac wrapper using six retained checkpoints on Amazon Reviews 2014 Beauty: 12,068 items, 22,363 test queries, beam 20 and history limit 210. Training was not rerun through this wrapper. ST5 uses 768-dimensional Sentence-T5 features; OpenAI uses 3072-dimensional `text-embedding-3-large` features. Values are mean ± sample SD across joint-training seeds 2025/2026/2027 with a fixed seed-2024 initializer; metrics are fractions. -| Features | Recall@5 | NDCG@5 | Recall@10 | NDCG@10 | Recall@20 | NDCG@20 | -| --- | --- | --- | --- | --- | --- | --- | -| ST5 | 0.041378 ± 0.001322 | 0.027401 ± 0.000342 | 0.063021 ± 0.001201 | 0.034355 ± 0.000252 | 0.093294 ± 0.003556 | 0.042015 ± 0.000854 | -| OpenAI | 0.044135 ± 0.001538 | 0.028961 ± 0.000630 | 0.069236 ± 0.001704 | 0.037026 ± 0.000674 | 0.105338 ± 0.001701 | 0.046100 ± 0.000367 | - -The retained features, initializers and checkpoints are not distributed with the -package; the presets alone do not reproduce this evaluation. - -### TIGER and RPG reference results - -For context, the following Beauty results are reported in the +For reference, the table also includes Beauty results from the [TIGER](../tiger/README.md#best-tiger-configs-vs-tuned-sasrec-beauty) and -[RPG](../rpg/README.md#graph-versus-exact-scoring) model documentation: +[RPG](../rpg/README.md#graph-versus-exact-scoring) documentation. DIGER rows +report three-seed mean ± sample SD; reference rows report a single run (seed 123). | Model / recipe | Features | Recall@5 | NDCG@5 | Recall@10 | NDCG@10 | Recall@20 | NDCG@20 | | --- | --- | --- | --- | --- | --- | --- | --- | -| TIGER / `PAISCHER_CONFIG`, RQVAE | ST5-base | 0.0419 | 0.0270 | 0.0634 | 0.0338 | 0.0953 | 0.0419 | -| TIGER / `PAISCHER_CONFIG`, rkmeans | OpenAI | 0.0405 | 0.0265 | 0.0678 | 0.0353 | 0.1045 | 0.0445 | -| RPG / `RPG_BEAUTY_CONFIG`, graph | OpenAI | 0.0533 | 0.0372 | 0.0789 | 0.0454 | 0.1113 | 0.0536 | +| DIGER | ST5 | 0.041378 ± 0.001322 | 0.027401 ± 0.000342 | 0.063021 ± 0.001201 | 0.034355 ± 0.000252 | 0.093294 ± 0.003556 | 0.042015 ± 0.000854 | +| [TIGER / `PAISCHER_CONFIG`, RQVAE](../tiger/README.md#best-tiger-configs-vs-tuned-sasrec-beauty) | ST5-base | 0.0419 | 0.0270 | 0.0634 | 0.0338 | 0.0953 | 0.0419 | +| DIGER | OpenAI | 0.044135 ± 0.001538 | 0.028961 ± 0.000630 | 0.069236 ± 0.001704 | 0.037026 ± 0.000674 | 0.105338 ± 0.001701 | 0.046100 ± 0.000367 | +| [TIGER / `PAISCHER_CONFIG`, rkmeans](../tiger/README.md#best-tiger-configs-vs-tuned-sasrec-beauty) | OpenAI | 0.0405 | 0.0265 | 0.0678 | 0.0353 | 0.1045 | 0.0445 | +| [RPG / `RPG_BEAUTY_CONFIG`, graph](../rpg/README.md#graph-versus-exact-scoring) | OpenAI | 0.0533 | 0.0372 | 0.0789 | 0.0454 | 0.1113 | 0.0536 | -These reference runs use seed 123; DIGER reports three-seed means above. TIGER uses history length 20 and constrained beam 50; RPG uses history length 50 and graph beam 20, versus DIGER's history limit 210 and unconstrained beam 20. All use Amazon Reviews 2014 Beauty, but identical filtering, feature arrays -and evaluation populations have not been established. The rows provide context -rather than a controlled comparison. +and evaluation populations have not been established. These are reference +results rather than a controlled comparison. + +For each DIGER feature setting, all six learning-rate candidates were evaluated, +and the configuration with the highest validation NDCG@10 was selected. This +limited search did not exhaustively explore model, tokenizer or decoding +settings. + +The retained DIGER features, initializers and checkpoints are not distributed +with the package; the presets alone do not reproduce this evaluation. ## Semantic ID evaluation