text_plain_output_1.png
from pathlib import Path

from torch.utils.data import Dataset, ConcatDataset

import pandas as pd

import torch

TRAIN_DATASET = '../input/starpredict-dataset/train.parquet'
VAL_DATASET = '../input/starpredict-dataset/val.parquet'
TEST_DATASET = '../input/starpredict-dataset/test.parquet'
SAMPLE_DATASET = '../input/starpredict-dataset/sample.parquet'

class YelpDataset(Dataset):
    def __init__(self, path: Path):
        super().__init__()
        self.data = pd.read_parquet(path)
    def __getitem__(self, key):
        row = self.data.iloc[key]
        return {'user_id': row['user_id_encode'], 'product_id': row['business_id_encode'], 'input_ids': torch.tensor(row['input_ids']), 'attention_mask': torch.tensor(row['attention_mask']), 'stars': row['stars_transform']}
    def __len__(self):
        return self.data.shape[0]
train_dataset = YelpDataset(TRAIN_DATASET)
val_dataset = YelpDataset(VAL_DATASET)
test_dataset = YelpDataset(TEST_DATASET)
sample_dataset = YelpDataset(SAMPLE_DATASET)

BATCH_SIZE = 8
LEARNING_RATE = 0.0003
EPOCHS = 3
MERGE_SIZE = 128
NUM_CLASSES = 5
NUM_WORKERS = 2

datasets = [train_dataset, val_dataset, test_dataset]
num_users = max((dataset.data['user_id_encode'].max() for dataset in datasets))
num_products = max((dataset.data['business_id_encode'].max() for dataset in datasets))
model = StarPredictSystem(num_users=num_users, num_products=num_products, merge_size=MERGE_SIZE, lr=LEARNING_RATE, num_classes=NUM_CLASSES, from_='bert-base-uncased')