utils.py

import os
import json
import torch
import random
import numpy as np
import pyarrow.parquet as pq
import torch.utils.data as data_utils
import torch.backends.cudnn as cudnn
from torch.utils.data.distributed import DistributedSampler


class FeatureDataset(torch.utils.data.Dataset):
    def __init__(self, args, data, dataset_type=None):
        # self.args = args
        self.data = data
        self.dataset_type = dataset_type
        self.max_len = args.max_len

    def __len__(self):
        return len(self.data["user_id"])

    def __getitem__(self, index):
        # [1 - N-1]
        history = self.data["history"][index][-self.max_len :]
        history = torch.tensor(history)
        history_len = torch.tensor(history.size(0))
        history = self.padding(history, max_len=self.max_len, padding_side="left")

        # [N]
        target = self.data["target"][index]
        target = torch.tensor(target)

        user_id = torch.tensor(self.data["user_id"][index])

        return history, target, history_len, user_id

    def padding(self, rows, max_len, padding_side="left"):
        cur_len = rows.size(0)
        pad_num = max(0, max_len - cur_len)

        if pad_num > 0:
            padding_tensor = torch.zeros(
                (pad_num), dtype=torch.long, device=rows.device
            )
            if padding_side == "left":
                rows = torch.cat([padding_tensor, rows], dim=0)
            elif padding_side == "right":
                rows = torch.cat([rows, padding_tensor], dim=0)
            else:
                raise ValueError("Invalid padding_side. Choose 'left' or 'right'.")

        return rows


def get_ordered_feature_description(args):
    r"""
    Load feature information from item_feature_explain.json and update args.
    All feature-related data (e.g., args.feature_description) will follow the same order as args.feature.

    Example:
        item_feature_explain.json:
            [
                {"id_num": 200, "col_index": 0, "name": "store_id"},
                {"id_num": 100, "col_index": 1, "name": "item_id"},
                {"id_num": 400, "col_index": 2, "name": "cate2_id"},
                {"id_num": 300, "col_index": 3, "name": "cate1_id"},
            ]
        args.feature:
            - item_id
            - store_id
            - cate1_id
            - cate2_id

        get_ordered_feature_description(args)

        args.feature_description:
            [
                {"id_num": 100, "col_index": 1, "name": "item_id"},
                {"id_num": 200, "col_index": 0, "name": "store_id"},
                {"id_num": 300, "col_index": 3, "name": "cate1_id"},
                {"id_num": 400, "col_index": 2, "name": "cate2_id"},
            ]
    """

    if getattr(args, "feature_mapping", None):
        # get semantic id as feature
        feature_path = os.path.join(args.feature_mapping, "item_feature_explain.json")
    else:
        # get normal feature
        feature_path = os.path.join(args.dataset, "item_feature_explain.json")

    with open(feature_path, "r", encoding="utf-8") as file:
        # list[dict]
        description = json.load(file)

    # dict[str, dict]
    # key("item_id") -> value({"id_num": 100, "col_index": 0, "name": "item_id"})
    desc_map = {field["name"]: field for field in description}

    # list[dict]
    # reorder to match args.feature order
    filtered_description = [desc_map[name] for name in args.feature if name in desc_map]

    args.feature_description = filtered_description

    return args


def remap_semantic_ids_to_unique(sids, F_list=None):
    r"""
    Remap semantic IDs in each feature column to global unique IDs.
    This ensures that IDs from different feature columns do not conflict,
    which is important when using embeddings in models like T5.

    Example:
        sids = tensor([[1, 2, 1],
                       [3, 1, 2],
                       [1, 3, 1]])

        remap_semantic_ids_to_unique(sids)

        sids = tensor([[1, 5, 7],
                       [3, 4, 8],
                       [1, 6, 7]])
    """

    if F_list is None:
        F = torch.max(sids, dim=0).values.to(dtype=torch.long)
    else:
        F = torch.tensor(F_list, dtype=torch.long, device=sids.device)

    offsets = torch.cat(
        (
            torch.tensor([0], dtype=torch.long, device=sids.device),
            torch.cumsum(F, dim=0)[:-1],
        )
    )
    offsets = offsets.unsqueeze(0)  # (1, C)
    mask = (sids != 0).long()
    new = sids + offsets * mask
    return new


def load_features_and_dataloaders(args):
    r"""
    args.item_feature_map (shape=[num_items+1, num_features]) stores the feature information for each item, indexed by item_id (start from 1).
    Each row corresponds to one item:
        args.item_feature_map[item_id] = [item_id, feature1, feature2, ...]
    """

    data_path = args.dataset
    needed_columns = args.feature

    if "item_id" not in args.feature:
        needed_columns = ["item_id"] + needed_columns

    if getattr(args, "feature_mapping", None):
        # get semantic id as feature
        df = pq.read_table(
            os.path.join(args.feature_mapping, "item_feature"),
            columns=needed_columns,
        ).to_pandas()
    else:
        # get normal feature
        df = pq.read_table(
            os.path.join(data_path, "item_feature"), columns=needed_columns
        ).to_pandas()

    # NOTE:
    # all feature indices start at 1, 0 (padding)
    # df["item_id"].max() = args.feature_description["item_id"]["id_num"] - 1
    # df["feature1"].max() = args.feature_description["feature1"]["id_num"] - 1
    # df["feature2"].max() = args.feature_description["feature2"]["id_num"] - 1
    max_item_id = df["item_id"].max()
    num_features = len(args.feature)
    full_tensor = torch.zeros((max_item_id + 1, num_features), dtype=torch.long)

    item_ids = df["item_id"].values
    features = df[args.feature].values
    full_tensor[item_ids] = torch.tensor(features, dtype=torch.long)

    if getattr(args, "train_type", None) == "t5":
        full_tensor = remap_semantic_ids_to_unique(full_tensor)

    args.item_feature_map = full_tensor.to(args.device)

    data_loaders = {}
    for split in ["train", "valid", "test"]:
        split_data = pq.read_table(os.path.join(data_path, split)).to_pydict()

        dataset = FeatureDataset(args, split_data, split)

        batch_size = args.batch_size
        if split in ["valid", "test"]:
            batch_size = getattr(args, "eval_batch_size", args.batch_size)

        if split == "train":
            args.num_users = max(split_data["user_id"]) + 1
            print(f"Number of users: {args.num_users}")

        if args.use_ddp:
            sampler = DistributedSampler(
                dataset,
                shuffle=(split == "train"),
                num_replicas=args.world_size,
                rank=args.rank,
            )
            # recommend: batch_size // args.world_size
            data_loader = data_utils.DataLoader(
                dataset,
                batch_size=batch_size,
                sampler=sampler,
                num_workers=4,
            )
        else:
            data_loader = data_utils.DataLoader(
                dataset,
                batch_size=batch_size,
                shuffle=True if split == "train" else False,
                num_workers=4,
                pin_memory=True,
            )

        data_loaders[split] = data_loader

    return data_loaders["train"], data_loaders["valid"], data_loaders["test"]


def fix_random_seed_as(random_seed):
    random.seed(random_seed)
    torch.manual_seed(random_seed)
    torch.cuda.manual_seed_all(random_seed)
    np.random.seed(random_seed)
    cudnn.deterministic = True
    cudnn.benchmark = False


def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group["lr"]