[{'role': 'system', 'content': "You are an expert in machine learning and your task is to classify the code snippets into categories from the provided list of categories. List of categories: ['helper_functions', 'load_data', 'data_exploration', 'data_preprocessing', 'evaluation', 'modelling', 'prediction', 'result_visualization', 'save_results', 'comment_only']"}, {'role': 'user', 'content': "Identify to which category the provided code belongs and assign the code snippet to at least 1 but up to 5 categories based on the probabilities, ignore if the probability of the category is negligible. Strictly provide your response as a list. Avoid providing any additional information, commentary, or personal opinions.\\nPython Code:class modified_dataset(torch.utils.data.Dataset):\n    def __init__(self, df, tokenizer):\n        p_text = []\n        offsets = []\n        at_mask = []\n        self.y_lst = df[['A-coref', 'B-coref']].apply(lambda row: target(row), axis = 1)\n        for row in tqdm(range(len(df))):\n            tokens, offset = tokenize(insert_tag(df.iloc[row]), tokenizer)\n            bla = tokenizer.encode_plus(tokens, max_length = pad_len, pad_to_max_length = True, return_token_type_ids = False)\n            p_text.append(bla['input_ids'])\n            at_mask.append(bla['attention_mask'])\n            offsets.append(offset)\n        self.p_text = torch.tensor(p_text)\n        self.offsets = torch.tensor(offsets)\n        self.at_mask = torch.tensor(at_mask)\n        return \n    def __len__(self):\n        return len(self.p_text)\n    def __getitem__(self,item):\n        return self.p_text[item], self.y_lst[item], self.offsets[item], self.at_mask[item]\n\nclass modified_dataset_test(torch.utils.data.Dataset):\n    def __init__(self, df, tokenizer):\n        p_text = []\n        offsets = []\n        at_mask = []\n        for row in tqdm(range(len(df))):\n            tokens, offset = tokenize(insert_tag(df.iloc[row]), tokenizer)\n            bla = tokenizer.encode_plus(tokens, max_length = pad_len, pad_to_max_length = True, return_token_type_ids = False)\n            p_text.append(bla['input_ids'])\n            at_mask.append(bla['attention_mask'])\n            offsets.append(offset)\n        self.p_text = torch.tensor(p_text)\n        self.offsets = torch.tensor(offsets)\n        self.at_mask = torch.tensor(at_mask)\n        return  \n    def __len__(self):\n        return len(self.p_text)\n    def __getitem__(self,item):\n        return self.p_text[item], self.offsets[item], self.at_mask[item]\n \ndef collate_fun(batch):\n    tmp_lst = list(zip(*batch))\n    return torch.stack(tmp_lst[0], axis = 0), torch.tensor(tmp_lst[1]), torch.stack(tmp_lst[2], axis = 0), torch.stack(tmp_lst[3], axis = 0)\n\ndef collate_fun2(batch):\n    tmp_lst = list(zip(*batch))\n    return torch.stack(tmp_lst[0], axis = 0), torch.stack(tmp_lst[1], axis = 0), torch.stack(tmp_lst[2], axis = 0)\n\ntrain_loader = DataLoader(\n        modified_dataset(df_train, tokenizer),\n        batch_size=18,\n        collate_fn=collate_fun,\n        shuffle=True,\n        drop_last=True,\n        num_workers=2)\nval_loader = DataLoader(\n        modified_dataset(df_val, tokenizer),\n        batch_size=30,\n        collate_fn=collate_fun,\n        shuffle=False,\n        num_workers=2)\ntest_loader = DataLoader(\n        modified_dataset(df_test, tokenizer),\n        batch_size=30,\n        collate_fn=collate_fun,\n        shuffle=False,\n        num_workers=2)\ntest_2_loader = DataLoader(\n        modified_dataset_test(test_2, tokenizer),\n        batch_size=30,\n        collate_fn=collate_fun2,\n        shuffle=False,\n        num_workers=2)"}, {'role': 'assistant', 'content': "['load_data']"}]