Create src/data_loader.py

src/data_loader.py

@@ -0,0 +1,14 @@
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+
+def get_dataloader(config, tokenizer, split='train'):
+    dataset = load_dataset("code_search_net", "python", split=split)
+    
+    def tokenize_function(examples):
+        return tokenizer(examples['whole_func_string'], truncation=True, padding='max_length', max_length=config['model']['max_length'])
+
+    tokenized_dataset = dataset.map(tokenize_function, batched=True)
+    tokenized_dataset = tokenized_dataset.remove_columns(['repo', 'path', 'func_name', 'whole_func_string', 'language', 'func_code_string', 'func_code_tokens', 'func_documentation_string', 'func_documentation_tokens', 'split_name', 'func_code_url'])
+    tokenized_dataset.set_format("torch")
+
+    return DataLoader(tokenized_dataset, batch_size=config['training']['batch_size'], shuffle=(split == 'train'))