43 lines
1.2 KiB
Python
43 lines
1.2 KiB
Python
"""
|
|
This file contains the HumorDataset class.
|
|
"""
|
|
import torch
|
|
import numpy as np
|
|
|
|
class HumorDataset(torch.utils.data.Dataset):
|
|
def __init__(self, data, labels, vocab_size=0, emb_dim=None):
|
|
self.original_indices = labels.index.to_list()
|
|
|
|
self.data = data
|
|
self.labels = labels.reset_index(drop=True)
|
|
self.vocab_size = vocab_size
|
|
self.emb_dim = emb_dim
|
|
|
|
# TODO: bug fix
|
|
self.shape = self.get_shape()
|
|
|
|
|
|
def __getitem__(self, idx):
|
|
item = {'input_ids': torch.tensor(self.data[idx], dtype=torch.float)}
|
|
item['labels'] = torch.tensor(self.labels[idx], dtype=torch.float)
|
|
return item
|
|
|
|
def __len__(self):
|
|
return len(self.labels)
|
|
|
|
def get_single_shape(self, data):
|
|
shape_data = None
|
|
if type(data) == list:
|
|
shape_data = len(data[0])
|
|
elif type(data) == torch.Tensor:
|
|
shape_data = data[0].shape
|
|
elif type(data) == np.ndarray:
|
|
shape_data = data[0].shape
|
|
return shape_data
|
|
|
|
def get_shape(self):
|
|
shape_data = self.get_single_shape(self.data)
|
|
shape_labels = self.get_single_shape(self.labels)
|
|
return shape_data, shape_labels
|
|
|