dataset.py

import json
import torch
from torch.utils.data import Dataset

class JsonlDataset(Dataset):
    def __init__(self, file_path):
        self.data = []
        with open(file_path, 'r') as f:
            lines = f.readlines()
            for idx, line in enumerate(lines):
                dict_line = json.loads(line)
                if 'idx' not in dict_line.keys():
                    dict_line['idx'] = idx
                else:
                    dict_line['idx'] = dict_line['idx']
                self.data.append(dict_line)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx]

def my_collate_fn(batch):
    batched_data = {}
    for key in batch[0]:
        value = batch[0][key]
        if type(value) == int and key != 'difficulty':
            batched_data[key] = torch.tensor([sample[key] for sample in batch])
        else:
            batched_data[key] = [sample[key] for sample in batch]
    return batched_data