From 61b0db3ff22e2c37e2540cd2bc093f3f4c408e09 Mon Sep 17 00:00:00 2001 From: hieuddo Date: Sun, 12 Jul 2026 23:50:09 +0800 Subject: [PATCH 1/2] load text for amazon review datasets --- cornac/datasets/README.md | 9 +- cornac/datasets/amazon_review.py | 149 ++++++++++++++++++-- tests/cornac/datasets/test_amazon_review.py | 81 +++++++++++ 3 files changed, 225 insertions(+), 14 deletions(-) diff --git a/cornac/datasets/README.md b/cornac/datasets/README.md index 9d1787b3..ed7aa4a1 100644 --- a/cornac/datasets/README.md +++ b/cornac/datasets/README.md @@ -296,10 +296,17 @@ Each user's reviews form one chronologically-ordered sequence. Interactions are | Amazon Sports (`sports`) | 35,598 | 18,357 | 296,337 | INT [1,5] | | Amazon Toys (`toys`) | 19,412 | 11,924 | 167,597 | INT [1,5] | +Item content text (title, price, brand, categories -- the features embedded with Sentence-T5 in the TIGER paper) is available via `amazon_review.load_text(category=...)`, built from the public product metadata and covering exactly the 5-core items. Passing `include_description=True` appends each item's product description to its text (cached separately); Paischer et al. (arXiv:2412.08604) found this beneficial for the Toys dataset, while attribute-only text works better for Beauty/Sports. + ```Python +from cornac.data import FeatureModality from cornac.datasets import amazon_review from cornac.eval_methods import NextItemEvaluation data = amazon_review.load_feedback(category="beauty") # UIRT tuples, chronological per user -eval_method = NextItemEvaluation.leave_last_out(data, fmt="UIRT") +texts, item_ids = amazon_review.load_text(category="beauty") # item content text, aligned ids +features = ... # embed texts, e.g. SentenceTransformer("sentence-t5-base").encode(texts) +eval_method = NextItemEvaluation.leave_last_out( + data, fmt="UIRT", item_feature=FeatureModality(features=features, ids=item_ids) +) ``` diff --git a/cornac/datasets/amazon_review.py b/cornac/datasets/amazon_review.py index 487e765c..7a8ee10a 100644 --- a/cornac/datasets/amazon_review.py +++ b/cornac/datasets/amazon_review.py @@ -21,6 +21,8 @@ Source: https://cseweb.ucsd.edu/~jmcauley/datasets/amazon/links.html """ +import ast +import csv import gzip import json import os @@ -65,6 +67,138 @@ def _preprocess(gz_path: str, csv_path: str) -> None: f.write(f"{user},{item},{rating},{timestamp}\n") +def _validate(category: str, version: str) -> None: + if category not in _CATEGORY_FILES: + raise ValueError(f"category='{category}' not supported; " f"choose one of {sorted(_CATEGORY_FILES)}") + if version != "2014": + raise ValueError(f"version='{version}' not supported; only '2014' (McAuley 5-core) " "is available") + + +def _reviews_csv(category: str, version: str) -> str: + stem = _CATEGORY_FILES[category] + gz_path = cache( + url=f"{_BASE_URL}/reviews_{stem}_5.json.gz", + relative_path=f"amazon_review/{category}_{version}.json.gz", + ) + csv_path = f"{gz_path[:-len('.json.gz')]}.csv" + if not os.path.exists(csv_path): + _preprocess(gz_path, csv_path) + return csv_path + + +def _item_text(meta: dict, include_description: bool = False) -> str: + """Flatten item metadata into one text string (title, price, brand, + categories -- the content features used by TIGER). + + When ``include_description`` is set, the item's ``description`` field (if + present and non-empty) is appended as a last ``Description: ...`` part. + """ + parts = [] + title = meta.get("title") + if title: + parts.append(f"Title: {title}") + price = meta.get("price") + if price is not None: + parts.append(f"Price: {price}") + brand = meta.get("brand") + if brand: + parts.append(f"Brand: {brand}") + categories = meta.get("categories") + if categories: + flat = [] + for path in categories: + for cat in path: + if cat not in flat: + flat.append(cat) + parts.append("Categories: " + ", ".join(flat)) + if include_description: + description = meta.get("description") + if description: + parts.append(f"Description: {description}") + return ". ".join(" ".join(part.split()) for part in parts) + + +def _preprocess_meta(meta_gz_path: str, reviews_csv_path: str, out_path: str, include_description: bool = False) -> None: + """Extract one text string per 5-core item from the raw category metadata. + + The 2014 metadata files contain Python dict literals (not valid JSON), + hence ``ast.literal_eval``. Items without a metadata entry get an empty + string so the output covers every item in the reviews file. + """ + keep = {} # item id -> insertion order (dict preserves order) + with open(reviews_csv_path) as f: + for line in f: + item = line.split(",")[1] + if item not in keep: + keep[item] = len(keep) + + texts = {} + with gzip.open(meta_gz_path, "rt", encoding="utf-8") as f: + for line in f: + meta = ast.literal_eval(line) + asin = meta.get("asin") + if asin in keep: + texts[asin] = _item_text(meta, include_description) + + with open(out_path, "w", newline="") as f: + writer = csv.writer(f) + for item in keep: + writer.writerow([item, texts.get(item, "")]) + + +def load_text(category: str, version: str = "2014", include_description: bool = False) -> (List, List): + """Load item content text (title, price, brand, categories) per item. + + Texts are built from the public product metadata of the same corpus as + :func:`load_feedback` and cover exactly the items appearing in the 5-core + reviews (items without a metadata entry get an empty string). These are + the content features embedded with Sentence-T5 in the TIGER paper. + + Parameters + ---------- + category: str, required + One of ``'beauty'``, ``'sports'``, ``'toys'``. + + version: str, default: '2014' + Dataset version. Only ``'2014'`` is currently supported. + + include_description: bool, default: False + If True, append each item's ``description`` field as a last + ``Description: ...`` part of its text. Paischer et al. + (arXiv:2412.08604) found this beneficial for the Toys dataset, while + attribute-only text works better for Beauty/Sports. Descriptions are + kept in full (downstream sentence encoders truncate as needed). The + description variant is cached separately (``*_text_desc.csv``) so it + never overwrites the attribute-only cache. + + Returns + ------- + texts: List + List of text documents, one per item. + + ids: List + List of item ids aligned with indices in `texts`. + """ + _validate(category, version) + reviews_csv_path = _reviews_csv(category, version) + suffix = "_text_desc" if include_description else "_text" + text_path = f"{reviews_csv_path[:-len('.csv')]}{suffix}.csv" + if not os.path.exists(text_path): + stem = _CATEGORY_FILES[category] + meta_gz_path = cache( + url=f"{_BASE_URL}/meta_{stem}.json.gz", + relative_path=f"amazon_review/meta_{category}_{version}.json.gz", + ) + _preprocess_meta(meta_gz_path, reviews_csv_path, text_path, include_description) + + texts, ids = [], [] + with open(text_path, newline="") as f: + for item, text in csv.reader(f): + ids.append(item) + texts.append(text) + return texts, ids + + def load_feedback(category: str, version: str = "2014", fmt: str = "UIRT", reader: Reader = None) -> List: """Load the user-item review feedback, chronologically ordered per user. @@ -90,19 +224,8 @@ def load_feedback(category: str, version: str = "2014", fmt: str = "UIRT", reade data: array-like Data in the form of a list of tuples (user, item, rating, timestamp). """ - if category not in _CATEGORY_FILES: - raise ValueError(f"category='{category}' not supported; " f"choose one of {sorted(_CATEGORY_FILES)}") - if version != "2014": - raise ValueError(f"version='{version}' not supported; only '2014' (McAuley 5-core) " "is available") - - stem = _CATEGORY_FILES[category] - gz_path = cache( - url=f"{_BASE_URL}/reviews_{stem}_5.json.gz", - relative_path=f"amazon_review/{category}_{version}.json.gz", - ) - csv_path = f"{gz_path[:-len('.json.gz')]}.csv" - if not os.path.exists(csv_path): - _preprocess(gz_path, csv_path) + _validate(category, version) + csv_path = _reviews_csv(category, version) reader = Reader() if reader is None else reader return reader.read(csv_path, fmt=fmt, sep=",") diff --git a/tests/cornac/datasets/test_amazon_review.py b/tests/cornac/datasets/test_amazon_review.py index 0166ed3d..6bf9c1bd 100644 --- a/tests/cornac/datasets/test_amazon_review.py +++ b/tests/cornac/datasets/test_amazon_review.py @@ -60,6 +60,87 @@ def test_load_feedback(self): self.assertGreater(len(data), 0) self.assertEqual(len(data[0]), 4) # (user, item, rating, timestamp) + def test_item_text(self): + meta = { + "asin": "iA", + "title": " Long Wig ", + "price": 11.83, + "brand": "Generic", + "categories": [["Beauty", "Hair Care"], ["Beauty", "Wigs"]], + } + self.assertEqual( + amazon_review._item_text(meta), + "Title: Long Wig. Price: 11.83. Brand: Generic. " + "Categories: Beauty, Hair Care, Wigs", + ) + # missing fields are skipped, not rendered empty + self.assertEqual(amazon_review._item_text({"title": "X"}), "Title: X") + + def test_item_text_with_description(self): + meta = { + "asin": "iA", + "title": "Wig", + "description": " A long soft wig. ", + } + # description is appended last and whitespace-normalized + self.assertEqual( + amazon_review._item_text(meta, include_description=True), + "Title: Wig. Description: A long soft wig.", + ) + # missing/empty description is skipped + self.assertEqual( + amazon_review._item_text({"title": "Wig"}, include_description=True), + "Title: Wig", + ) + self.assertEqual( + amazon_review._item_text({"title": "Wig", "description": ""}, include_description=True), + "Title: Wig", + ) + # include_description=False is unchanged even when description present + self.assertEqual(amazon_review._item_text(meta), "Title: Wig") + + def test_preprocess_meta_covers_all_review_items(self): + # Metadata lines are Python dict literals (single quotes), only items + # in the reviews file are kept, and items without metadata get "". + with tempfile.TemporaryDirectory() as d: + reviews_csv = os.path.join(d, "reviews.csv") + with open(reviews_csv, "w") as f: + f.write("u1,iA,5.0,100\nu1,iB,4.0,200\nu2,iA,3.0,150\n") + + meta_gz = os.path.join(d, "meta.json.gz") + with gzip.open(meta_gz, "wt", encoding="utf-8") as f: + f.write("{'asin': 'iA', 'title': 'Item A', 'price': 9.99}\n") + f.write("{'asin': 'iZ', 'title': 'Not reviewed'}\n") + + out_path = os.path.join(d, "text.csv") + amazon_review._preprocess_meta(meta_gz, reviews_csv, out_path) + + with open(out_path) as f: + rows = [ln.rstrip("\n").split(",", 1) for ln in f] + + self.assertEqual(rows[0], ["iA", "Title: Item A. Price: 9.99"]) + self.assertEqual(rows[1], ["iB", ""]) # covered, but no metadata + self.assertEqual(len(rows), 2) # iZ excluded + + def test_preprocess_meta_include_description(self): + # include_description plumbs through to the emitted text. + with tempfile.TemporaryDirectory() as d: + reviews_csv = os.path.join(d, "reviews.csv") + with open(reviews_csv, "w") as f: + f.write("u1,iA,5.0,100\n") + + meta_gz = os.path.join(d, "meta.json.gz") + with gzip.open(meta_gz, "wt", encoding="utf-8") as f: + f.write("{'asin': 'iA', 'title': 'Item A', 'description': 'Nice item'}\n") + + out_path = os.path.join(d, "text.csv") + amazon_review._preprocess_meta(meta_gz, reviews_csv, out_path, include_description=True) + + with open(out_path) as f: + rows = [ln.rstrip("\n").split(",", 1) for ln in f] + + self.assertEqual(rows[0], ["iA", "Title: Item A. Description: Nice item"]) + if __name__ == "__main__": unittest.main() From 3831b5cd3549a7231c44f4fdf7d0e34740bdf126 Mon Sep 17 00:00:00 2001 From: hieuddo Date: Wed, 15 Jul 2026 02:29:24 +0800 Subject: [PATCH 2/2] add TIGER model --- README.md | 3 +- cornac/models/__init__.py | 1 + cornac/models/tiger/__init__.py | 18 + cornac/models/tiger/grid_config.py | 52 +++ cornac/models/tiger/paischer_config.py | 67 +++ cornac/models/tiger/recom_tiger.py | 593 +++++++++++++++++++++++++ cornac/models/tiger/requirements.txt | 2 + cornac/models/tiger/tiger.py | 289 ++++++++++++ examples/README.md | 2 + examples/tiger_example.py | 87 ++++ 10 files changed, 1113 insertions(+), 1 deletion(-) create mode 100644 cornac/models/tiger/__init__.py create mode 100644 cornac/models/tiger/grid_config.py create mode 100644 cornac/models/tiger/paischer_config.py create mode 100644 cornac/models/tiger/recom_tiger.py create mode 100644 cornac/models/tiger/requirements.txt create mode 100644 cornac/models/tiger/tiger.py create mode 100644 examples/tiger_example.py diff --git a/README.md b/README.md index 18cb586f..a9d61e86 100644 --- a/README.md +++ b/README.md @@ -152,7 +152,8 @@ The table below lists the recommendation models/algorithms featured in Cornac. E | :--: | --------------- | :--: | :---------: | :-----: | | 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_cu116.txt), CPU / GPU | [quick-start](https://github.com/PreferredAI/HypAR) -| 2023 | [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) +| 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) | 2021 | [Bilateral Variational Autoencoder for Collaborative Filtering (BiVAECF)](cornac/models/bivaecf), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.bivaecf.recom_bivaecf), [paper](https://dl.acm.org/doi/pdf/10.1145/3437963.3441759) | Collaborative Filtering / Content-Based | [requirements](cornac/models/bivaecf/requirements.txt), CPU / GPU | [quick-start](https://github.com/PreferredAI/bi-vae), [deep-dive](https://github.com/recommenders-team/recommenders/blob/main/examples/02_model_collaborative_filtering/cornac_bivae_deep_dive.ipynb) | | [Transformers4Rec-style Unified Transformer Recommender (TransformerRec)](cornac/models/transformer_rec), [docs](https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.transformer_rec.recom_transformer_rec), [paper](https://dl.acm.org/doi/10.1145/3460231.3474255) | Next-Item | [requirements](cornac/models/transformer_rec/requirements.txt), CPU / GPU | [quick-start](examples/transformer_rec_diginetica.py) diff --git a/cornac/models/__init__.py b/cornac/models/__init__.py index b20c7494..1e9e3e5a 100644 --- a/cornac/models/__init__.py +++ b/cornac/models/__init__.py @@ -83,6 +83,7 @@ from .spop import SPop from .svd import SVD from .tifuknn import TIFUKNN +from .tiger import TIGER from .transformer_rec import TransformerRec from .trirank import TriRank from .upcf import UPCF diff --git a/cornac/models/tiger/__init__.py b/cornac/models/tiger/__init__.py new file mode 100644 index 00000000..7222d965 --- /dev/null +++ b/cornac/models/tiger/__init__.py @@ -0,0 +1,18 @@ +# 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 .grid_config import GRID_CONFIG +from .paischer_config import PAISCHER_CONFIG +from .recom_tiger import TIGER diff --git a/cornac/models/tiger/grid_config.py b/cornac/models/tiger/grid_config.py new file mode 100644 index 00000000..cb6009a1 --- /dev/null +++ b/cornac/models/tiger/grid_config.py @@ -0,0 +1,52 @@ +# 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. +# ============================================================================ +"""GRID handbook recipe for TIGER (Ju et al., 2025, arXiv 2507.22224). + +Trades some accuracy for much faster training: the residual k-means tokenizer +needs no gradient training at all (GRID Table 1 finds it *outperforms* RQ-VAE +anyway), and the seq2seq stage runs a short epoch budget with best-on-val +checkpoint selection instead of a long fixed schedule. + +Usage:: + + from cornac.models.tiger import GRID_CONFIG, TIGER + + model = TIGER(**{**GRID_CONFIG, "seed": 123}) + +Verbatim from the GRID paper: optimizer settings (Adam lr 5e-4, weight decay +1e-6, batch 256, constant lr), NDCG@10 validation, rkmeans tokenizer with the +paper's (3 levels, 256 codes) shape and the paper's 4+4-layer transformer +(both cornac defaults, so not repeated here). Adapted to cornac's epoch-based +trainer: GRID counts optimizer *steps* (one step = one minibatch update, i.e. +``n_steps = n_epochs * ceil(n_train_samples / batch_size)``, where a session +of length T contributes T-1 samples) and early-stops on step-based validation +intervals (every 100 steps, patience 10); cornac instead trains ``n_epochs`` +and keeps the best-on-val checkpoint, so ``n_epochs``/``val_eval_every``/ +``val_sample`` here are pragmatic equivalents, not GRID values. +""" + +GRID_CONFIG = dict( + tokenizer="rkmeans", + learning_rate=5e-4, + weight_decay=1e-6, + batch_size=256, + lr_schedule="constant", + model_selection="best", + val_metric="ndcg", + val_k=10, + n_epochs=50, + val_eval_every=1, + val_sample=2000, +) diff --git a/cornac/models/tiger/paischer_config.py b/cornac/models/tiger/paischer_config.py new file mode 100644 index 00000000..14bb93d4 --- /dev/null +++ b/cornac/models/tiger/paischer_config.py @@ -0,0 +1,67 @@ +# 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. +# ============================================================================ +"""Paischer et al. recipe for TIGER (arXiv 2412.08604) - the best documented +reproduction of the original paper's numbers (matches Beauty/Sports). + +Taken from ``configs/setting/MenderTok_Beauty.yaml`` in +https://github.com/facebookresearch/preference_discerning (their TIGER +baseline shares the trainer). Slower than :data:`GRID_CONFIG` but tuned for +accuracy. Their item embeddings are Sentence-T5-XXL on title/price/brand/ +categories text (descriptions instead for Toys) - the embedding side lives in +the experiment script, not here. + +Usage:: + + from cornac.models.tiger import PAISCHER_CONFIG, TIGER + + model = TIGER(**{**PAISCHER_CONFIG, "seed": 123}) + +Verbatim from their config: RQ-VAE (hidden 768/512/256, latent 128, AdamW wd +0.1, lr 1e-3, batch 2048, 8000 epochs, standardized inputs), seq2seq (6+6 +layers, dropout 0.2, lr 3e-4 cosine with 10k warmup steps, weight decay +0.035, batch 64), beam width 30. Adapted to cornac's epoch-based trainer: +they cap training at 200k optimizer *steps* (one step = one minibatch update, +so ``n_steps = n_epochs * ceil(n_train_samples / batch_size)``, where a +session of length T contributes T-1 samples) with step-based early stopping +(patience 15). At batch 64 on Amazon-Beauty-scale data (~128k training +samples, ~2k steps/epoch), ``n_epochs=100`` lands near their 200k-step cap, +with best-on-val checkpoint selection instead of early stopping +(``val_eval_every``/``val_sample`` are pragmatic choices, not theirs). +""" + +PAISCHER_CONFIG = dict( + tokenizer="rqvae", + feature_standardize=True, + rqvae_hidden_dims=(768, 512, 256), + rqvae_latent_dim=128, + rqvae_n_epochs=8000, + rqvae_batch_size=2048, + rqvae_weight_decay=0.1, + num_enc_layers=6, + num_dec_layers=6, + dropout=0.2, + learning_rate=3e-4, + lr_schedule="cosine", + warmup_steps=10000, + weight_decay=0.035, + batch_size=64, + n_beams=30, + model_selection="best", + val_metric="ndcg", + val_k=10, + n_epochs=100, + val_eval_every=5, + val_sample=2000, +) diff --git a/cornac/models/tiger/recom_tiger.py b/cornac/models/tiger/recom_tiger.py new file mode 100644 index 00000000..d4a68051 --- /dev/null +++ b/cornac/models/tiger/recom_tiger.py @@ -0,0 +1,593 @@ +# 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 +from collections import defaultdict + +import numpy as np +from tqdm.auto import trange + +from cornac.models.recommender import NextItemRecommender + +from ...utils import get_rng +from ..seq_utils import session_seq_iter + +SUPPORTED_SCORING = ("beam", "exact") +SUPPORTED_TOKENIZERS = ("rqvae", "rkmeans") +SUPPORTED_LR_SCHEDULES = ("constant", "cosine") + + +class TIGER(NextItemRecommender): + """TIGER: Recommender Systems with Generative Retrieval. + + Two-stage generative retrieval with semantic IDs: (1) an RQ-VAE quantizes + precomputed item content embeddings into hierarchical semantic IDs + (``rqvae_num_levels`` codebooks of ``rqvae_codebook_size`` codes, plus an + extra level that disambiguates collisions); (2) a T5-style encoder-decoder + is trained to generate the next item's semantic ID from the session + history's semantic-ID tokens. + + Item content embeddings must be provided through the evaluation method, + e.g.:: + + NextItemEvaluation.from_splits( + ..., item_feature=FeatureModality(features=embs, ids=item_ids) + ) + + where ``embs`` are precomputed text/content embeddings (e.g., from + sentence-transformers) covering every known item. + + Two ready-made configurations ship with the model: + :data:`~cornac.models.tiger.GRID_CONFIG` (GRID handbook recipe — fast, + no tokenizer training) and :data:`~cornac.models.tiger.PAISCHER_CONFIG` + (Paischer et al. recipe — best documented reproduction accuracy), e.g. + ``TIGER(**{**PAISCHER_CONFIG, "seed": 123})``. + + Parameters + ---------- + name: string, default: 'TIGER' + The name of the recommender model. + + tokenizer: str, optional, default: 'rqvae' + Semantic-ID tokenizer. 'rqvae' trains the residual-quantized VAE + (paper-faithful); 'rkmeans' runs residual k-means directly on the item + features (level-by-level k-means on the residuals, no encoder/decoder + and no gradient training), the simpler GRID-handbook baseline. Both + share the collision-disambiguation level and the seq2seq stage. + + feature_standardize: bool, optional, default: False + When True, z-score the item features per dimension (over items) before + tokenizing. Paischer et al. standardize the RQ-VAE inputs. + + rqvae_latent_dim: int, optional, default: 32 + RQ-VAE latent (codeword) dimension. + + rqvae_hidden_dims: tuple of int, optional, default: (512, 256, 128) + Hidden layer sizes of the RQ-VAE encoder (decoder mirrors them). + + rqvae_num_levels: int, optional, default: 3 + Number of residual codebooks (semantic-ID levels before the + collision-disambiguation level). + + rqvae_codebook_size: int, optional, default: 256 + Number of codes per codebook. + + rqvae_beta: float, optional, default: 0.25 + Commitment loss weight of the RQ-VAE. + + rqvae_n_epochs: int, optional, default: 200 + rqvae_learning_rate: float, optional, default: 0.001 + rqvae_batch_size: int, optional, default: 1024 + RQ-VAE training settings. + + rqvae_weight_decay: float, optional, default: 0.0 + AdamW weight decay for RQ-VAE training (Paischer et al. use 0.1). The + default 0.0 reproduces plain Adam. + + d_model: int, optional, default: 128 + d_ff: int, optional, default: 1024 + num_heads: int, optional, default: 6 + d_kv: int, optional, default: 64 + num_enc_layers: int, optional, default: 4 + num_dec_layers: int, optional, default: 4 + dropout: float, optional, default: 0.1 + Transformer (T5-style) architecture settings, defaults per the paper + (~13M parameters). + + max_len: int, optional, default: 20 + Maximum number of history items fed to the encoder. The encoder input + is ``max_len * (rqvae_num_levels + 1)`` tokens. + + n_epochs: int, optional, default: 20 + learning_rate: float, optional, default: 0.001 + weight_decay: float, optional, default: 0.0001 + batch_size: int, optional, default: 256 + Seq2seq training settings. Note: the paper uses Adagrad and an + inverse-sqrt schedule; we use Adam/AdamW with constant lr (as in the + GRID framework), which converges much faster at these model sizes. + + lr_schedule: str, optional, default: 'constant' + Seq2seq learning-rate schedule. 'constant' keeps ``learning_rate`` + fixed; 'cosine' does linear warmup over ``warmup_steps`` steps then + cosine decay to ~0 over the remaining total training steps + (steps_per_epoch * n_epochs), as in Paischer et al. + + warmup_steps: int, optional, default: 10000 + Number of linear-warmup steps when ``lr_schedule='cosine'``. + + model_selection: str, optional, default: 'last' + One of 'last' or 'best'. When 'best' and a ``val_set`` is given, the + seq2seq weights with the highest validation score (evaluated every + ``val_eval_every`` epochs on up to ``val_sample`` val sessions) are + restored at the end of ``fit``. The tokenizer is fixed before seq2seq + training, so only the seq2seq model is snapshotted. + + val_metric: str, optional, default: 'ndcg' + val_eval_every: int, optional, default: 5 + val_k: int, optional, default: 10 + Metric, epoch cadence, and cutoff K used for best-on-val selection. + During validation ``n_beams`` is raised to at least ``val_k``. + + val_sample: int, optional, default: 2000 + Maximum number of val sessions scored per evaluation (None = all). + Beam-search scoring is expensive, so a fixed deterministic subsample + (drawn once with the model's rng) is reused across epochs. + + scoring: str, optional, default: 'beam' + 'beam' (paper-faithful constrained beam search; only the top + ``n_beams`` items receive real scores, so set ``n_beams`` >= the + largest metric cutoff K) or 'exact' (teacher-forced log-likelihood of + every item's semantic ID; exact full ranking, slower per user). + + n_beams: int, optional, default: 20 + Beam width for scoring='beam'. + + scoring_batch_size: int, optional, default: 2048 + Item chunk size for scoring='exact'. + + device: str, optional, default: 'auto' + 'auto' selects 'cuda' if available, otherwise 'cpu'. + + trainable: bool, optional, default: True + When False, the model will not be re-trained. + + verbose: bool, optional, default: False + When True, running logs are displayed. + + seed: int, optional, default: None + Random seed for weight initialization. + + References + ---------- + Rajput, S. et al. (2023). Recommender Systems with Generative Retrieval. + NeurIPS. https://arxiv.org/pdf/2305.05065 + + Ju, C. et al. (2025). Generative Recommendation with Semantic IDs: A + Practitioner's Handbook (GRID). https://github.com/snap-research/GRID + https://arxiv.org/abs/2507.22224 + + Paischer, F. et al. (2024). Preference Discerning with LLM-Enhanced + Generative Retrieval. https://arxiv.org/abs/2412.08604 + """ + + def __init__( + self, + name="TIGER", + tokenizer="rqvae", + feature_standardize=False, + rqvae_latent_dim=32, + rqvae_hidden_dims=(512, 256, 128), + rqvae_num_levels=3, + rqvae_codebook_size=256, + rqvae_beta=0.25, + rqvae_n_epochs=200, + rqvae_learning_rate=0.001, + rqvae_batch_size=1024, + rqvae_weight_decay=0.0, + d_model=128, + d_ff=1024, + num_heads=6, + d_kv=64, + num_enc_layers=4, + num_dec_layers=4, + dropout=0.1, + max_len=20, + n_epochs=20, + learning_rate=0.001, + weight_decay=0.0001, + batch_size=256, + lr_schedule="constant", + warmup_steps=10000, + model_selection="last", + val_metric="ndcg", + val_eval_every=5, + val_k=10, + val_sample=2000, + scoring="beam", + n_beams=20, + scoring_batch_size=2048, + device="auto", + trainable=True, + verbose=False, + seed=None, + ): + super().__init__(name, trainable=trainable, verbose=verbose) + if scoring not in SUPPORTED_SCORING: + raise ValueError(f"scoring='{scoring}' not supported; choose from {SUPPORTED_SCORING}") + if tokenizer not in SUPPORTED_TOKENIZERS: + raise ValueError(f"tokenizer='{tokenizer}' not supported; choose from {SUPPORTED_TOKENIZERS}") + if lr_schedule not in SUPPORTED_LR_SCHEDULES: + raise ValueError(f"lr_schedule='{lr_schedule}' not supported; choose from {SUPPORTED_LR_SCHEDULES}") + if model_selection not in ("last", "best"): + raise ValueError(f"model_selection='{model_selection}' not supported; choose 'last' or 'best'") + self.tokenizer = tokenizer + self.feature_standardize = feature_standardize + self.rqvae_latent_dim = rqvae_latent_dim + self.rqvae_hidden_dims = rqvae_hidden_dims + self.rqvae_num_levels = rqvae_num_levels + self.rqvae_codebook_size = rqvae_codebook_size + self.rqvae_beta = rqvae_beta + self.rqvae_n_epochs = rqvae_n_epochs + self.rqvae_learning_rate = rqvae_learning_rate + self.rqvae_batch_size = rqvae_batch_size + self.rqvae_weight_decay = rqvae_weight_decay + self.d_model = d_model + self.d_ff = d_ff + self.num_heads = num_heads + self.d_kv = d_kv + self.num_enc_layers = num_enc_layers + self.num_dec_layers = num_dec_layers + self.dropout = dropout + self.max_len = max_len + self.n_epochs = n_epochs + self.learning_rate = learning_rate + self.weight_decay = weight_decay + self.batch_size = batch_size + self.lr_schedule = lr_schedule + self.warmup_steps = warmup_steps + self.model_selection = model_selection + self.val_metric = val_metric + self.val_eval_every = val_eval_every + self.val_k = val_k + self.val_sample = val_sample + self.scoring = scoring + self.n_beams = n_beams + self.scoring_batch_size = scoring_batch_size + self.device = device + self.seed = seed + self.rng = get_rng(seed) + + def _get_item_features(self): + item_feature = getattr(self.train_set, "item_feature", None) + features = getattr(item_feature, "features", None) + if features is None: + raise ValueError( + "TIGER requires precomputed item content embeddings. Provide them " + "via NextItemEvaluation.from_splits(..., item_feature=" + "FeatureModality(features=..., ids=...))." + ) + if features.shape[0] < self.total_items: + raise ValueError( + f"item_feature has {features.shape[0]} rows but {self.total_items} " + "items are known; every item (train/val/test) needs a feature vector." + ) + return np.asarray(features[: self.total_items], dtype="float32") + + def _fit_rqvae(self, torch, feats_t): + from .tiger import RQVAE + + self.rqvae = RQVAE( + 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, + beta=self.rqvae_beta, + ).to(self.device_) + self.rqvae.kmeans_init_codebooks(feats_t) + opt = torch.optim.AdamW( + self.rqvae.parameters(), + lr=self.rqvae_learning_rate, + weight_decay=self.rqvae_weight_decay, + ) + + n = feats_t.size(0) + progress_bar = trange(1, self.rqvae_n_epochs + 1, disable=not self.verbose, desc="RQ-VAE") + for _ in progress_bar: + self.rqvae.train() + used = torch.zeros( + self.rqvae_num_levels, + self.rqvae_codebook_size, + dtype=torch.bool, + device=self.device_, + ) + perm = torch.randperm(n, device=self.device_) + total_loss, cnt = 0.0, 0 + for start in range(0, n, self.rqvae_batch_size): + batch = feats_t[perm[start : start + self.rqvae_batch_size]] + ids, _, loss_recon, loss_rq = self.rqvae(batch) + loss = loss_recon + loss_rq + opt.zero_grad() + loss.backward() + opt.step() + for level in range(self.rqvae_num_levels): + used[level, ids[:, level]] = True + total_loss += loss.item() * len(batch) + cnt += len(batch) + sample = feats_t[perm[: min(n, 8192)]] + self.rqvae.restart_dead_codes(sample, used) + progress_bar.set_postfix(loss=(total_loss / cnt)) + + def _fit_rkmeans(self, torch, feats_t): + """Residual k-means tokenizer (GRID handbook): level-by-level k-means on + the item features, subtracting the assigned centroid before the next + level. Returns (N, num_levels) int64 codes; centroids are cached (numpy) + for pickling.""" + from .tiger import _kmeans + + self.rkmeans_centroids = [] + codes = [] + r = feats_t + for _ in range(self.rqvae_num_levels): + centroids = _kmeans(r, self.rqvae_codebook_size) + ids = torch.cdist(r, centroids).argmin(dim=1) + r = r - centroids[ids] + self.rkmeans_centroids.append(centroids.cpu().numpy()) + codes.append(ids.cpu().numpy()) + return np.stack(codes, axis=1).astype("int64") + + def _tokenize(self, torch, feats_t): + if self.tokenizer == "rkmeans": + return self._fit_rkmeans(torch, feats_t) + self._fit_rqvae(torch, feats_t) + self.rqvae.eval() + return np.concatenate( + [ + self.rqvae.encode(feats_t[start : start + self.rqvae_batch_size]).cpu().numpy() + for start in range(0, feats_t.size(0), self.rqvae_batch_size) + ] + ).astype("int64") + + def _build_semantic_ids(self, codes): + # extra level disambiguating items that share the same code tuple + counters = defaultdict(int) + dedup = np.zeros(len(codes), dtype="int64") + for i, row in enumerate(map(tuple, codes)): + dedup[i] = counters[row] + counters[row] += 1 + self.sid_table = np.concatenate([codes, dedup[:, None]], axis=1) + self.level_sizes = [self.rqvae_codebook_size] * self.rqvae_num_levels + [int(dedup.max()) + 1] + + # prefix trie for constrained beam search + sid -> item lookup + children = [defaultdict(set) for _ in self.level_sizes] + self.sid_to_item = {} + for i, row in enumerate(self.sid_table): + sid = tuple(int(v) for v in row) + for level in range(len(sid)): + children[level][sid[:level]].add(sid[level]) + self.sid_to_item[sid] = i + self.prefix_children = [ + {prefix: np.fromiter(sorted(tokens), dtype="int64") for prefix, tokens in level_children.items()} + for level_children in children + ] + if self.verbose: + n_collisions = int((dedup > 0).sum()) + print( + f"Semantic IDs assigned: {len(self.sid_table)} items, " + f"{n_collisions} collisions, dedup level size {self.level_sizes[-1]}" + ) + + def _fit_seq2seq(self, torch, val_set): + from .tiger import TIGERSeq2Seq + + self.model = TIGERSeq2Seq( + 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, + ).to(self.device_) + + # per-item encoder tokens (offset, 0 = pad); extra all-pad row for pad_idx + self.pad_idx = self.total_items + offsets = self.model.offsets.cpu().numpy() + self.enc_token_table = np.zeros((self.total_items + 1, len(self.level_sizes)), dtype="int64") + self.enc_token_table[: self.total_items] = self.sid_table + offsets + + opt = torch.optim.AdamW( + self.model.parameters(), + lr=self.learning_rate, + weight_decay=self.weight_decay, + ) + scheduler = self._make_lr_scheduler(torch, opt) + + best_state, best_val = None, -float("inf") + select_best = self.model_selection == "best" and val_set is not None + val_sessions = self._val_sessions(val_set) if select_best else None + val_metric = self._make_val_metric() if select_best else None + + progress_bar = trange(1, self.n_epochs + 1, disable=not self.verbose, desc="TIGER") + for epoch_id in progress_bar: + self.model.train() + total_loss, cnt = 0.0, 0 + for inc, (_, hist_iids, out_iids) in enumerate( + session_seq_iter( + self.train_set, + pad_index=self.pad_idx, + batch_size=self.batch_size, + max_len=self.max_len, + n_sample=0, + rng=self.rng, + shuffle=True, + ) + ): + enc_tokens = self.enc_token_table[hist_iids].reshape(len(hist_iids), -1) + enc_tokens_t = torch.tensor(enc_tokens, dtype=torch.long, device=self.device_) + enc_mask_t = (enc_tokens_t != 0).float() + target_t = torch.tensor(self.sid_table[out_iids], dtype=torch.long, device=self.device_) + opt.zero_grad() + loss = self.model(enc_tokens_t, enc_mask_t, target_t) + loss.backward() + opt.step() + if scheduler is not None: + scheduler.step() + total_loss += loss.item() * len(hist_iids) + cnt += len(hist_iids) + if inc % 10 == 0 and cnt > 0: + progress_bar.set_postfix(loss=(total_loss / cnt)) + + if select_best and epoch_id % self.val_eval_every == 0: + score = self._validate(val_sessions, val_metric) + if score > best_val: + best_val = score + best_state = {n: p.detach().clone() for n, p in self.model.state_dict().items()} + + if best_state is not None: + self.model.load_state_dict(best_state) + + def _make_lr_scheduler(self, torch, opt): + if self.lr_schedule != "cosine": + return None + n_samples = sum(len(m) - 1 for m in self.train_set.sessions.values() if len(m) >= 2) + steps_per_epoch = max(1, math.ceil(n_samples / self.batch_size)) + total_steps = max(1, steps_per_epoch * self.n_epochs) + + def lr_lambda(step): + if step < self.warmup_steps: + return (step + 1) / max(1, self.warmup_steps) + progress = (step - self.warmup_steps) / max(1, total_steps - self.warmup_steps) + return 0.5 * (1.0 + math.cos(math.pi * min(1.0, progress))) + + return torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda) + + def _make_val_metric(self): + from ...metrics import AUC, MRR, NDCG, Recall + + name = self.val_metric.lower() + if name == "recall": + return Recall(k=self.val_k) + if name == "ndcg": + return NDCG(k=self.val_k) + if name == "auc": + return AUC() + if name == "mrr": + return MRR() + raise ValueError(f"val_metric='{self.val_metric}' not supported; choose from recall/ndcg/auc/mrr") + + def _val_sessions(self, val_set): + """Collect (user_idx, session_items) for last-item eval, deterministically + subsampled to ``val_sample`` sessions (fixed across epochs).""" + sessions = [] + for [_], [mapped_ids], [session_items] in val_set.si_iter(batch_size=1, shuffle=False): + if len(session_items) < 2: + continue + user_idx = int(val_set.uir_tuple[0][mapped_ids[0]]) + sessions.append((user_idx, [int(i) for i in session_items])) + if self.val_sample is not None and len(sessions) > self.val_sample: + idx = self.rng.choice(len(sessions), size=self.val_sample, replace=False) + sessions = [sessions[i] for i in sorted(idx)] + return sessions + + def _validate(self, val_sessions, metric): + """Mean ``metric`` over ``val_sessions`` (mode 'last'), imitating + ranking_eval. Beam scoring only fills the top ``n_beams`` items, so + ``n_beams`` is temporarily raised to at least ``val_k``.""" + num_items = self.train_set.num_items + orig_beams, self.n_beams = self.n_beams, max(self.n_beams, self.val_k) + item_indices = np.arange(num_items) + results = [] + for user_idx, session_items in val_sessions: + target = session_items[-1] + if target >= num_items: + continue + gt_pos = np.array([target]) + gt_neg = np.delete(item_indices, target) + item_rank, item_scores = self.rank(user_idx, item_indices, history_items=session_items[:-1]) + results.append( + metric.compute( + gt_pos=gt_pos, + gt_neg=gt_neg, + pd_rank=item_rank, + pd_scores=item_scores, + item_indices=item_indices, + ) + ) + self.n_beams = orig_beams + return float(np.mean(results)) if results else 0.0 + + def fit(self, train_set, val_set=None): + super().fit(train_set, val_set) + if not self.trainable: + return self + + import torch + + torch.manual_seed(self.seed if self.seed is not None else 0) + self.device_ = ("cuda" if torch.cuda.is_available() else "cpu") if self.device == "auto" else self.device + + feats = self._get_item_features() + if self.feature_standardize: + mean = feats.mean(axis=0) + std = feats.std(axis=0) + std[std == 0] = 1.0 + feats = ((feats - mean) / std).astype("float32") + feats_t = torch.tensor(feats, device=self.device_) + + codes = self._tokenize(torch, feats_t) + self._build_semantic_ids(codes) + self._fit_seq2seq(torch, val_set) + + # keep pickles portable across GPU/CPU boxes; moved back in score() + self.model.to("cpu").eval() + if self.tokenizer == "rqvae": + self.rqvae.to("cpu").eval() + return self + + def _ensure_device(self, torch): + if self.device_ == "cuda" and not torch.cuda.is_available(): + self.device_ = "cpu" + if next(self.model.parameters()).device.type != torch.device(self.device_).type: + self.model.to(self.device_) + + 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) + hist = list(history_items)[-self.max_len :] + hist = [self.pad_idx] * (self.max_len - len(hist)) + hist + enc_tokens_t = torch.tensor( + self.enc_token_table[hist].reshape(1, -1), + dtype=torch.long, + device=self.device_, + ) + enc_mask_t = (enc_tokens_t != 0).float() + self.model.eval() + with torch.no_grad(): + if self.scoring == "beam": + beams, logps = self.model.generate_beam(enc_tokens_t, enc_mask_t, self.n_beams, self.prefix_children) + scores = np.full(self.total_items, -1e10, dtype="float") + for sid, lp in zip(beams, logps): + scores[self.sid_to_item[sid]] = lp + else: # "exact" + sid_table_t = torch.tensor(self.sid_table, dtype=torch.long, device=self.device_) + scores = self.model.score_all_items( + enc_tokens_t, enc_mask_t, sid_table_t, self.scoring_batch_size + ).astype("float") + return scores diff --git a/cornac/models/tiger/requirements.txt b/cornac/models/tiger/requirements.txt new file mode 100644 index 00000000..a808c3f6 --- /dev/null +++ b/cornac/models/tiger/requirements.txt @@ -0,0 +1,2 @@ +torch>=1.12.0 +transformers>=4.30.0 diff --git a/cornac/models/tiger/tiger.py b/cornac/models/tiger/tiger.py new file mode 100644 index 00000000..31303a6d --- /dev/null +++ b/cornac/models/tiger/tiger.py @@ -0,0 +1,289 @@ +# 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 TIGER (Rajput et al., 2023). + +``RQVAE`` quantizes item content embeddings into hierarchical semantic IDs; +``TIGERSeq2Seq`` is a small T5-style encoder-decoder trained to generate the +next item's semantic ID from the session history's semantic-ID tokens. +""" + +import copy + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import T5Config +from transformers.models.t5.modeling_t5 import T5Stack + + +def _mlp(input_dim, hidden_dims, output_dim): + dims = [input_dim, *hidden_dims, output_dim] + layers = [] + for i in range(len(dims) - 1): + layers.append(nn.Linear(dims[i], dims[i + 1])) + if i < len(dims) - 2: + layers.append(nn.ReLU()) + return nn.Sequential(*layers) + + +@torch.no_grad() +def _kmeans(x, k, n_iters=10): + """K-means++ seeding followed by Lloyd iterations. Returns (k, dim) centroids.""" + n = x.size(0) + if n < k: + idx = torch.randint(0, n, (k,), device=x.device) + return x[idx] + 1e-4 * torch.randn(k, x.size(1), device=x.device, dtype=x.dtype) + centroids = torch.empty(k, x.size(1), device=x.device, dtype=x.dtype) + centroids[0] = x[torch.randint(0, n, (1,), device=x.device)] + d2 = torch.cdist(x, centroids[:1]).squeeze(1).pow_(2) + for i in range(1, k): + centroids[i] = x[torch.multinomial(d2 + 1e-12, 1)] + d2 = torch.minimum(d2, torch.cdist(x, centroids[i : i + 1]).squeeze(1).pow_(2)) + for _ in range(n_iters): + ids = torch.cdist(x, centroids).argmin(dim=1) + sums = torch.zeros_like(centroids) + sums.index_add_(0, ids, x) + counts = torch.bincount(ids, minlength=k) + empty = counts == 0 + centroids = sums / counts.clamp_min(1).unsqueeze(1) + if empty.any(): + centroids[empty] = x[torch.randint(0, n, (int(empty.sum()),), device=x.device)] + return centroids + + +class RQVAE(nn.Module): + """Residual-Quantized VAE over item content embeddings (TIGER paper Sec. 4.1). + + Encoder MLP -> ``num_levels`` residual codebooks (straight-through + estimator, VQ-VAE codebook + commitment losses) -> mirror decoder MLP + with MSE reconstruction loss. + """ + + def __init__( + self, + input_dim, + hidden_dims=(512, 256, 128), + latent_dim=32, + num_levels=3, + codebook_size=256, + beta=0.25, + ): + super().__init__() + self.num_levels = num_levels + self.codebook_size = codebook_size + self.beta = beta + self.encoder = _mlp(input_dim, hidden_dims, latent_dim) + self.decoder = _mlp(latent_dim, tuple(reversed(hidden_dims)), input_dim) + self.codebooks = nn.Parameter( + torch.randn(num_levels, codebook_size, latent_dim) * 0.01 + ) + + def _quantize(self, z): + ids = [] + q = torch.zeros_like(z) + loss_rq = z.new_zeros(()) + r = z + for level in range(self.num_levels): + level_ids = torch.cdist(r, self.codebooks[level]).argmin(dim=1) + e = self.codebooks[level][level_ids] + loss_rq = loss_rq + F.mse_loss(e, r.detach()) + self.beta * F.mse_loss(r, e.detach()) + ids.append(level_ids) + q = q + e + # detach so each codebook is only updated by its own level's loss + r = r - e.detach() + return torch.stack(ids, dim=1), q, loss_rq + + def forward(self, x): + z = self.encoder(x) + ids, q, loss_rq = self._quantize(z) + z_q = z + (q - z).detach() # straight-through estimator + x_hat = self.decoder(z_q) + loss_recon = F.mse_loss(x_hat, x) + return ids, x_hat, loss_recon, loss_rq + + @torch.no_grad() + def encode(self, x): + """Assign semantic-ID codes. Returns LongTensor of shape (B, num_levels).""" + ids, _, _ = self._quantize(self.encoder(x)) + return ids + + @torch.no_grad() + def kmeans_init_codebooks(self, x, n_iters=10): + """Initialize each codebook with k-means on the encoder residuals, + level by level, to prevent early codebook collapse.""" + r = self.encoder(x) + for level in range(self.num_levels): + centroids = _kmeans(r, self.codebook_size, n_iters=n_iters) + self.codebooks[level].copy_(centroids) + ids = torch.cdist(r, centroids).argmin(dim=1) + r = r - centroids[ids] + + @torch.no_grad() + def restart_dead_codes(self, x, used): + """Reassign codes unused during the last epoch (``used``: bool (L, K)) + to random encoder residuals from ``x``. Returns #codes restarted.""" + r = self.encoder(x) + n_restarted = 0 + for level in range(self.num_levels): + dead = ~used[level] + if dead.any(): + idx = torch.randint(0, r.size(0), (int(dead.sum()),), device=r.device) + self.codebooks[level][dead] = r[idx] + n_restarted += int(dead.sum()) + ids = torch.cdist(r, self.codebooks[level]).argmin(dim=1) + r = r - self.codebooks[level][ids] + return n_restarted + + +class TIGERSeq2Seq(nn.Module): + """T5-style encoder-decoder over semantic-ID tokens (TIGER paper Fig. 2.b). + + A single embedding table holds all levels' tokens with cumulative offsets + (index 0 reserved for padding) and is shared between encoder and decoder; + each level has its own output head. A learned BOS embedding prompts the + decoder. + """ + + 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, + ): + super().__init__() + self.level_sizes = [int(s) for s in level_sizes] + self.num_levels = len(self.level_sizes) + # token 0 = padding; level l tokens occupy [offsets[l], offsets[l] + level_sizes[l]) + offsets = np.concatenate(([1], 1 + np.cumsum(self.level_sizes[:-1]))) + self.register_buffer("offsets", torch.as_tensor(offsets, dtype=torch.long)) + self.token_emb = nn.Embedding(1 + sum(self.level_sizes), d_model, padding_idx=0) + + cfg = T5Config( + vocab_size=1, # internal embed_tokens is unused; we always pass inputs_embeds + d_model=d_model, + d_ff=d_ff, + d_kv=d_kv, + num_heads=num_heads, + num_layers=num_enc_layers, + dropout_rate=dropout, + is_encoder_decoder=False, + use_cache=False, + ) + enc_cfg = copy.deepcopy(cfg) + enc_cfg.is_decoder = False + dec_cfg = copy.deepcopy(cfg) + dec_cfg.is_decoder = True + dec_cfg.num_layers = num_dec_layers + # T5Stack is semi-internal; if it breaks in a future transformers + # release, fall back to T5ForConditionalGeneration(cfg).encoder/.decoder + self.encoder = T5Stack(enc_cfg) + self.decoder = T5Stack(dec_cfg) + + self.bos = nn.Parameter(torch.randn(1, 1, d_model) * 0.02) + self.heads = nn.ModuleList( + [nn.Linear(d_model, s, bias=False) for s in self.level_sizes] + ) + + def encode_history(self, enc_tokens, enc_mask): + """enc_tokens: (B, S) offset tokens with 0 = pad; enc_mask: (B, S) float.""" + return self.encoder( + inputs_embeds=self.token_emb(enc_tokens), attention_mask=enc_mask + ).last_hidden_state + + def _decode(self, dec_sids, enc_out, enc_mask): + """Decoder pass on BOS + un-offset level tokens ``dec_sids`` (B, t) or None. + Returns hidden states (B, t+1, d_model).""" + inputs = self.bos.expand(enc_out.size(0), -1, -1) + if dec_sids is not None and dec_sids.size(1) > 0: + emb = self.token_emb(dec_sids + self.offsets[: dec_sids.size(1)]) + inputs = torch.cat([inputs, emb], dim=1) + return self.decoder( + inputs_embeds=inputs, + encoder_hidden_states=enc_out, + encoder_attention_mask=enc_mask, + use_cache=False, + ).last_hidden_state + + def forward(self, enc_tokens, enc_mask, target_sids): + """Teacher-forced training loss: sum over levels of cross-entropy. + target_sids: (B, num_levels) un-offset codes.""" + enc_out = self.encode_history(enc_tokens, enc_mask) + h = self._decode(target_sids[:, :-1], enc_out, enc_mask) # (B, num_levels, d) + loss = h.new_zeros(()) + for level in range(self.num_levels): + loss = loss + F.cross_entropy( + self.heads[level](h[:, level]), target_sids[:, level] + ) + return loss + + @torch.no_grad() + def generate_beam(self, enc_tokens, enc_mask, n_beams, prefix_children): + """Constrained beam search for a single history (batch size 1). + + At each level, candidate tokens are masked to ``prefix_children[level][prefix]`` + so every surviving beam is a valid item's semantic ID. The decoder is + recomputed per level (length <= num_levels + 1), avoiding KV-cache + bookkeeping. Returns (list of sid tuples, log-probs) sorted descending. + """ + enc_out = self.encode_history(enc_tokens, enc_mask) # (1, S, d) + beams = [()] + beam_lp = enc_out.new_zeros(1) + for level, size in enumerate(self.level_sizes): + n_b = len(beams) + dec_sids = ( + torch.tensor(beams, dtype=torch.long, device=enc_out.device) + if level > 0 + else None + ) + h = self._decode( + dec_sids, enc_out.expand(n_b, -1, -1), enc_mask.expand(n_b, -1) + ) + logp = F.log_softmax(self.heads[level](h[:, -1]), dim=-1) # (n_b, size) + allowed = torch.full_like(logp, float("-inf")) + for i, beam in enumerate(beams): + allowed[i, prefix_children[level][beam]] = 0.0 + total = (beam_lp.unsqueeze(1) + logp + allowed).flatten() + k = min(n_beams, int(torch.isfinite(total).sum())) + top = total.topk(k) + beams = [beams[j // size] + (j % size,) for j in top.indices.tolist()] + beam_lp = top.values + return beams, beam_lp.cpu().numpy() + + @torch.no_grad() + def score_all_items(self, enc_tokens, enc_mask, sid_table, batch_size): + """Exact teacher-forced log-likelihood of every item's semantic ID for a + single history (batch size 1). sid_table: (N, num_levels) un-offset codes. + Returns a 1-D numpy array of length N.""" + enc_out = self.encode_history(enc_tokens, enc_mask) # (1, S, d) + num_items = sid_table.size(0) + scores = enc_out.new_empty(num_items) + for start in range(0, num_items, batch_size): + target = sid_table[start : start + batch_size] + n = target.size(0) + h = self._decode( + target[:, :-1], enc_out.expand(n, -1, -1), enc_mask.expand(n, -1) + ) + s = h.new_zeros(n) + for level in range(self.num_levels): + logp = F.log_softmax(self.heads[level](h[:, level]), dim=-1) + s = s + logp.gather(1, target[:, level : level + 1]).squeeze(1) + scores[start : start + n] = s + return scores.cpu().numpy() diff --git a/examples/README.md b/examples/README.md index 7fc06d79..9aa1592d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -130,6 +130,8 @@ [transformer_rec_diginetica.py](transformer_rec_diginetica.py) - Example of Transformer-based Recommendation models (SASRec, TransformerRec with CLM/MLM objectives) with Diginetica dataset. +[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. + ---- ## Next-Basket Algorithms diff --git a/examples/tiger_example.py b/examples/tiger_example.py new file mode 100644 index 00000000..b6c40791 --- /dev/null +++ b/examples/tiger_example.py @@ -0,0 +1,87 @@ +# 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. +# ============================================================================ +"""TIGER (generative retrieval with semantic IDs) on Amazon Beauty (2014). + +Reproduces the TIGER paper setup: 5-core Amazon Beauty reviews with per-user +leave-last-out splitting (the paper's protocol; see +``NextItemEvaluation.from_timestamps`` for a leakage-free alternative), item +content text (title/price/brand/categories) embedded with Sentence-T5, and +the Paischer et al. training recipe shipped as +``cornac.models.tiger.PAISCHER_CONFIG`` -- the best documented reproduction +of the paper's numbers. + +Requires ``sentence-transformers`` on top of the model requirements +(torch, transformers). + +Expected results (test split, beam scoring): Recall@5 ~= 0.042, +NDCG@5 ~= 0.027, vs 0.0454 / 0.0321 reported in the paper -- on par with the +best published reproductions. Training takes about an hour on one GPU; +evaluation decodes a beam per test user and takes a comparable amount of time. +""" + +import torch +from sentence_transformers import SentenceTransformer + +import cornac +from cornac.data import FeatureModality +from cornac.datasets import amazon_review +from cornac.eval_methods import NextItemEvaluation +from cornac.metrics import MRR, NDCG, Recall +from cornac.models import TIGER +from cornac.models.tiger import PAISCHER_CONFIG + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +print(f"using device: {DEVICE}") + +data = amazon_review.load_feedback(category="beauty") +texts, item_ids = amazon_review.load_text(category="beauty") + +encoder = SentenceTransformer("sentence-t5-base", device=DEVICE) +features = encoder.encode(texts, batch_size=64, show_progress_bar=True) +del encoder # release encoder memory before training +if DEVICE == "cuda": + torch.cuda.empty_cache() + +next_item_eval = NextItemEvaluation.leave_last_out( + data=data, + exclude_unknowns=True, + verbose=True, + item_feature=FeatureModality(features=features, ids=item_ids), +) + +models = [ + TIGER( + **{ + **PAISCHER_CONFIG, + "device": DEVICE, + "verbose": True, + "seed": 123, + } + ), +] + +metrics = [ + Recall(k=5), + Recall(k=10), + NDCG(k=5), + NDCG(k=10), + MRR(), +] + +cornac.Experiment( + eval_method=next_item_eval, + models=models, + metrics=metrics, +).run()