application_vnd.jupyter.stderr_output_1.png
from pathlib import Path

from torch import nn

from torch.utils.data import Dataset, ConcatDataset

from torchmetrics.functional import accuracy, f1_score, precision, recall

import pandas as pd

import pytorch_lightning as pl

import torch

import torch_optimizer as optim

import transformers as T

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)

class StarPredictSystem(pl.LightningModule):
    def __init__(self, num_users, num_products, merge_size=512, lr=0.001, num_classes=5, from_='bert-base-uncased'):
        super().__init__()
        self.lr = lr
        self.num_classes = num_classes
        self.bert = T.DistilBertModel.from_pretrained('distilbert-base-uncased')
        self.bert_classifier = nn.Linear(self.bert.config.hidden_size, merge_size)
        self.user_embedding = nn.Embedding(num_users, merge_size)
        self.product_embedding = nn.Embedding(num_products, merge_size)
        self.classifier = nn.Sequential(nn.Linear(merge_size * 3, 64), nn.ReLU(), nn.Linear(64, num_classes), nn.Softmax(dim=1))
    def forward(self, x):
        user_x = x['user_id']
        user_x = self.user_embedding(user_x)
        product_x = x['product_id']
        product_x = self.product_embedding(product_x)
        text_x = self.bert(input_ids=x['input_ids'], attention_mask=x['attention_mask']).last_hidden_state[:, 0]
        text_x = self.bert_classifier(text_x)
        x = torch.cat([user_x, product_x, text_x], dim=-1)
        x = self.classifier(x)
        return x
    def configure_optimizers(self):
        return optim.Lamb(self.parameters(), lr=self.lr, weight_decay=0.02)
    def training_step(self, batch, batch_idx):
        y = batch['stars'].long()
        y_hat = self(batch)
        loss = F.cross_entropy(y_hat, y)
        acc = accuracy(y_hat, y)
        self.log('acc', acc, prog_bar=True, batch_size=batch['stars'].shape[0])
        return loss
    def validation_step(self, batch, batch_idx):
        y = batch['stars'].long()
        y_hat = self(batch)
        loss = F.cross_entropy(y_hat, y)
        metrics = {'val_loss': loss, 'val_acc': accuracy(y_hat, y), 'val_f1': f1_score(y_hat, y), 'val_prec': precision(y_hat, y), 'val_recall': recall(y_hat, y)}
        self.log_dict(metrics, batch_size=batch['stars'].shape[0])
        return metrics
    def test_step(self, batch, batch_idx):
        y = batch['stars'].long() - 1
        y_hat = self(batch)
        loss = F.cross_entropy(y_hat, y)
        metrics = {'test_loss': loss, 'test_acc': accuracy(y_hat, y), 'test_f1': f1_score(y_hat, y), 'test_prec': precision(y_hat, y), 'test_recall': recall(y_hat, y)}
        self.log_dict(metrics, batch_size=batch['stars'].shape[0])
        return metrics

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')

datamodule = pl.LightningDataModule.from_datasets(train_dataset=train_dataset, val_dataset=val_dataset, test_dataset=test_dataset, batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)
trainer = pl.Trainer(max_epochs=EPOCHS, gpus=-1)
trainer.fit(model, datamodule=datamodule)