-
Notifications
You must be signed in to change notification settings - Fork 171
Expand file tree
/
Copy pathtiger_example.py
More file actions
87 lines (75 loc) · 2.84 KB
/
Copy pathtiger_example.py
File metadata and controls
87 lines (75 loc) · 2.84 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
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()