From 2f7abfd2505b3e5923454995ed3d11bc3fe1b02b Mon Sep 17 00:00:00 2001 From: Iskander Bagautdinov Date: Sun, 14 Jun 2026 19:56:12 +0300 Subject: [PATCH] aligning sasrec implementation with CIKM2020-S3Rec --- notebooks/DatasetProcessing.ipynb | 137 ++++++++++++------------- tiger/configs/sasrec_train_config.json | 11 +- tiger/modeling/dataset/base.py | 2 +- tiger/modeling/dataset/samplers.py | 19 +++- tiger/modeling/models/base.py | 21 +--- tiger/modeling/models/sasrec.py | 29 +++--- tiger/train_sasrec.py | 4 +- 7 files changed, 107 insertions(+), 116 deletions(-) diff --git a/notebooks/DatasetProcessing.ipynb b/notebooks/DatasetProcessing.ipynb index 00bc16d..1b44d8b 100644 --- a/notebooks/DatasetProcessing.ipynb +++ b/notebooks/DatasetProcessing.ipynb @@ -2,10 +2,8 @@ "cells": [ { "cell_type": "code", - "execution_count": null, "id": "3bdb292f", "metadata": {}, - "outputs": [], "source": [ "from collections import defaultdict\n", "\n", @@ -21,69 +19,56 @@ "from torch.utils.data import DataLoader\n", "\n", "from tqdm import tqdm as tqdm" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "66d9b312", "metadata": {}, - "outputs": [], "source": [ "interactions_dataset_path = '../data/Beauty/Beauty_5.json'\n", "metadata_path = '../data/Beauty/metadata.json'\n", "\n", - "interactions_output_path = '../data/Beauty/inter_new.json'\n", - "embeddings_output_path = '../data/Beauty/content_embeddings.pkl'" - ] + "interactions_output_path = '../data/Beauty/inter.json'\n", + "embeddings_output_path = '../data/Beauty/content_embeddings.pkl'\n" + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "6ed4dffb", "metadata": {}, + "source": "df = defaultdict(list)\n\nwith open(interactions_dataset_path, 'r') as f:\n for line in f.readlines():\n review = json.loads(line)\n df['user_id'].append(review['reviewerID'])\n df['item_id'].append(review['asin'])\n df['timestamp'].append(review['unixReviewTime'])\n\nprint(f'Number of events: {len(df[\"user_id\"])}')\n\ndf = pl.from_dict(df).with_row_index('_row_idx')\n", "outputs": [], - "source": [ - "df = defaultdict(list)\n", - "\n", - "with open(interactions_dataset_path, 'r') as f:\n", - " for line in f.readlines():\n", - " review = json.loads(line)\n", - " df['user_id'].append(review['reviewerID'])\n", - " df['item_id'].append(review['asin'])\n", - " df['timestamp'].append(review['unixReviewTime'])\n", - "\n", - "print(f'Number of events: {len(df[\"user_id\"])}')\n", - "\n", - "df = pl.from_dict(df)" - ] + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "c26746c4", "metadata": {}, - "outputs": [], "source": [ "df.head()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "adcf5713", "metadata": {}, - "outputs": [], "source": [ "filtered_df = df.clone()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "c0bbf9ba", "metadata": {}, - "outputs": [], "source": [ "# Processing dataset to get core-5 state in case full dataset is provided\n", "is_changed = True\n", @@ -107,22 +92,24 @@ " )\n", "\n", " old_size = len(filtered_df)\n", + " filtered_df = filtered_df.join(\n", + " good_users, on='user_id', how='inner',\n", + " ).join(\n", + " good_items, on='item_id', how='inner',\n", + " )\n", "\n", - " new_df = filtered_df.join(good_users, on='user_id', how='inner')\n", - " new_df = new_df.join(good_items, on='item_id', how='inner')\n", - "\n", - " new_size = len(new_df)\n", + " if len(filtered_df) == old_size:\n", + " is_changed = False\n", "\n", - " filtered_df = new_df\n", - " is_changed = old_size != new_size\n" - ] + "filtered_df = filtered_df.sort('_row_idx').drop('_row_idx')\n" + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "218a9348", "metadata": {}, - "outputs": [], "source": [ "unique_values = filtered_df[\"user_id\"].unique(maintain_order=True).to_list()\n", "user_ids_mapping = {value: i for i, value in enumerate(unique_values)}\n", @@ -139,83 +126,83 @@ ")\n", "\n", "filtered_df.head()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "34604fe6", "metadata": {}, - "outputs": [], "source": [ "item_ids_mapping_df = pl.from_dict({\n", " 'old_item_id': list(item_ids_mapping.keys()),\n", " 'new_item_id': list(item_ids_mapping.values())\n", "})\n", "item_ids_mapping_df.head()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "6017e65c", "metadata": {}, - "outputs": [], "source": [ "filtered_df.head()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "9efd1983", "metadata": {}, - "outputs": [], "source": [ "filtered_df = filtered_df.sort([\"user_id\", \"timestamp\"])\n", "\n", "grouped_filtered_df = filtered_df.group_by(\"user_id\", maintain_order=True).agg(pl.all())" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "fd51c525", "metadata": {}, - "outputs": [], "source": [ "item_ids_mapping_df.head()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "8b0821da", "metadata": {}, - "outputs": [], "source": [ "grouped_filtered_df.head()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "dc222d59", "metadata": {}, - "outputs": [], "source": [ "print('Users count:', filtered_df.select('user_id').unique().shape[0])\n", "print('Items count:', filtered_df.select('item_id').unique().shape[0])\n", "print('Actions count:', filtered_df.shape[0])\n", "print('Avg user history len:', np.mean(list(map(lambda x: x[0], grouped_filtered_df.select(pl.col('item_id').list.len()).rows()))))" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "d07a2e91", "metadata": {}, - "outputs": [], "source": [ "json_data = {}\n", "for user_id, item_ids, _ in grouped_filtered_df.iter_rows():\n", @@ -223,7 +210,9 @@ "\n", "with open(interactions_output_path, 'w') as f:\n", " json.dump(json_data, f, indent=2)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", @@ -235,10 +224,8 @@ }, { "cell_type": "code", - "execution_count": null, "id": "6361c7a5", "metadata": {}, - "outputs": [], "source": [ "def getDF(path):\n", " i = 0\n", @@ -252,14 +239,14 @@ "\n", "df = getDF(metadata_path)\n", "df.head()" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "971fa89c", "metadata": {}, - "outputs": [], "source": [ "def preprocess(row: pd.Series):\n", " row = row.fillna(\"None\")\n", @@ -278,24 +265,24 @@ " filtered_df[\"combined_text\"] = filtered_df.apply(preprocess, axis=1)\n", "\n", " return filtered_df\n" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "3b0dd5d5", "metadata": {}, - "outputs": [], "source": [ "data = get_data(pl.from_pandas(df), item_ids_mapping)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "12e622ff", "metadata": {}, - "outputs": [], "source": [ "device = torch.device('cuda:6')\n", "\n", @@ -353,7 +340,9 @@ "\n", "with open(embeddings_output_path, 'wb') as f:\n", " pickle.dump(new_df, f)\n" - ] + ], + "outputs": [], + "execution_count": null } ], "metadata": { diff --git a/tiger/configs/sasrec_train_config.json b/tiger/configs/sasrec_train_config.json index dd319bb..e61dce7 100644 --- a/tiger/configs/sasrec_train_config.json +++ b/tiger/configs/sasrec_train_config.json @@ -2,7 +2,7 @@ "experiment_name": "sasrec_beauty", "dataset": { "inter_json_path": "../data/Beauty/inter.json", - "max_sequence_length": 20, + "max_sequence_length": 50, "sampler_type": "sasrec" }, "dataloader": { @@ -14,13 +14,14 @@ "num_heads": 2, "num_layers": 2, "dim_feedforward": 256, - "dropout": 0.3, - "activation": "relu", + "dropout": 0.5, + "activation": "gelu", "top_k": 20, - "layer_norm_eps": 1e-8, + "layer_norm_eps": 1e-12, "initializer_range": 0.02 }, "optimizer": { "lr": 0.001 - } + }, + "early_stopping_threshold": 10 } diff --git a/tiger/modeling/dataset/base.py b/tiger/modeling/dataset/base.py index 48b44b2..f6512e5 100644 --- a/tiger/modeling/dataset/base.py +++ b/tiger/modeling/dataset/base.py @@ -75,7 +75,7 @@ def create(cls, inter_json_path, max_sequence_length, sampler_type, is_extended= LOGGER.info(f'Test dataset size: {len(test_dataset)}') LOGGER.info(f'Max item id: {max_item_id}') - train_sampler = TrainSampler(train_dataset, sampler_type, max_sequence_length=max_sequence_length) + train_sampler = TrainSampler(train_dataset, sampler_type, max_sequence_length=max_sequence_length, num_items=max_item_id + 1) validation_sampler = EvalSampler(validation_dataset, max_sequence_length=max_sequence_length) test_sampler = EvalSampler(test_dataset, max_sequence_length=max_sequence_length) diff --git a/tiger/modeling/dataset/samplers.py b/tiger/modeling/dataset/samplers.py index 7ad0d36..262dd13 100644 --- a/tiger/modeling/dataset/samplers.py +++ b/tiger/modeling/dataset/samplers.py @@ -1,27 +1,42 @@ import copy +import random class TrainSampler: - def __init__(self, dataset, prediction_type, max_sequence_length): + def __init__(self, dataset, prediction_type, max_sequence_length, num_items=None): self._dataset = dataset self._prediction_type = prediction_type self._max_sequence_length = max_sequence_length + self._num_items = num_items self._transforms = { 'sasrec': self._all_items_transform, 'tiger': self._last_item_transform } + def _sample_negative(self, exclude: set) -> int: + assert self._num_items is not None, 'num_items required for negative sampling' + while True: + neg = random.randint(0, self._num_items - 1) + if neg not in exclude: + return neg + def _all_items_transform(self, sample): item_sequence = sample['item.ids'][-self._max_sequence_length:][:-1] next_item_sequence = sample['item.ids'][-self._max_sequence_length:][1:] + + exclude = set(sample['item.ids']) + negative_ids = [self._sample_negative(exclude) for _ in next_item_sequence] + return { 'user.ids': sample['user.ids'], 'user.length': len(sample['user.ids']), 'item.ids': item_sequence, 'item.length': len(item_sequence), 'labels.ids': next_item_sequence, - 'labels.length': len(next_item_sequence) + 'labels.length': len(next_item_sequence), + 'negative.ids': negative_ids, + 'negative.length': len(negative_ids), } def _last_item_transform(self, sample): diff --git a/tiger/modeling/models/base.py b/tiger/modeling/models/base.py index 5032a87..52347ab 100644 --- a/tiger/modeling/models/base.py +++ b/tiger/modeling/models/base.py @@ -10,27 +10,12 @@ def _init_weights(self, initializer_range): if 'norm' in key: nn.init.ones_(value.data) else: - nn.init.trunc_normal_( - value.data, - std=initializer_range, - a=-2 * initializer_range, - b=2 * initializer_range - ) + nn.init.normal_(value.data, mean=0.0, std=initializer_range) elif 'bias' in key: nn.init.zeros_(value.data) elif 'codebook' in key: - nn.init.trunc_normal_( - value.data, - std=initializer_range, - a=-2 * initializer_range, - b=2 * initializer_range - ) + nn.init.normal_(value.data, mean=0.0, std=initializer_range) elif 'bos_embedding' in key: - nn.init.trunc_normal_( - value.data, - std=initializer_range, - a=-2 * initializer_range, - b=2 * initializer_range, - ) + nn.init.normal_(value.data, mean=0.0, std=initializer_range) else: raise ValueError(f'Unknown transformer weight: {key}') \ No newline at end of file diff --git a/tiger/modeling/models/sasrec.py b/tiger/modeling/models/sasrec.py index d0c23b5..880e7b4 100644 --- a/tiger/modeling/models/sasrec.py +++ b/tiger/modeling/models/sasrec.py @@ -24,6 +24,7 @@ def __init__( self._num_items = num_items self._num_heads = num_heads self._embedding_dim = embedding_dim + self._max_sequence_length = max_sequence_length self._item_embeddings = nn.Embedding( num_embeddings=num_items, @@ -36,6 +37,9 @@ def __init__( self._topk_k = topk_k + self._embedding_layer_norm = nn.LayerNorm(embedding_dim, eps=layer_norm_eps) + self._embedding_dropout = nn.Dropout(dropout) + transformer_encoder_layer = nn.TransformerEncoderLayer( d_model=embedding_dim, nhead=num_heads, @@ -77,16 +81,13 @@ def forward(self, inputs): index=all_positive_sample_events[..., None] )[:, 0] # (all_batch_items) + negative_sample_events = inputs['negative.ids'] # (all_batch_events) + negative_scores = torch.gather( input=all_scores, dim=1, - index=torch.randint( - low=0, - high=all_scores.shape[1], - size=all_positive_sample_events.shape, - device=all_positive_sample_events.device - )[..., None] - )[:, 0] # (all_batch_items) + index=negative_sample_events[..., None] + )[:, 0] # (all_batch_events) return { 'positive_scores': positive_scores, @@ -124,7 +125,7 @@ def _apply_sequential_encoder(self, events, lengths): embeddings, mask = create_masked_tensor( data=embeddings, - lengths=lengths + lengths=lengths, ) # (batch_size, seq_len, embedding_dim), (batch_size, seq_len) batch_size = mask.shape[0] @@ -134,16 +135,16 @@ def _apply_sequential_encoder(self, events, lengths): start=0, end=seq_len, device=mask.device )[None].expand(batch_size, -1) # (batch_size, seq_len) - position_embeddings = self._position_embeddings(positions) # (batch_size, seq_len, embedding_dim) - position_embeddings[~mask] = 0 + position_embeddings = self._position_embeddings(positions) - embeddings = embeddings + position_embeddings # (batch_size, seq_len, embedding_dim) - embeddings[~mask] = 0 - causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len).bool().to(embeddings.device) # (seq_len, seq_len) + embeddings = embeddings + position_embeddings + embeddings = self._embedding_layer_norm(embeddings) + embeddings = self._embedding_dropout(embeddings) + causal_mask = nn.Transformer.generate_square_subsequent_mask(seq_len).bool().to(embeddings.device) embeddings = self._encoder( src=embeddings, mask=causal_mask, src_key_padding_mask=~mask - ) # (batch_size, seq_len, embedding_dim) + ) return embeddings, mask diff --git a/tiger/train_sasrec.py b/tiger/train_sasrec.py index c5a7f4a..64a9461 100644 --- a/tiger/train_sasrec.py +++ b/tiger/train_sasrec.py @@ -39,7 +39,7 @@ def main(): train_dataloader = DataLoader( dataset=train_sampler, batch_size=config['dataloader']['train_batch_size'], - drop_last=True, + drop_last=False, shuffle=True, collate_fn=batch_processor ) @@ -86,7 +86,7 @@ def main(): output_prefix='loss' ) - optimizer = torch.optim.AdamW( + optimizer = torch.optim.Adam( model.parameters(), lr=config['optimizer']['lr'], )