code: def create_dataset(data_dir, task):
      if task == "AdvTest":
          train_dataset = csn_search_train(data_dir, task, "train")
          val_dataset = advtest_search_eval_text(
              data_dir, task, "valid"
          )
          test_dataset = advtest_search_eval_text(
              data_dir, task, "test"
          )
          codebase_dataset = csn_search_eval_code(
              data_dir, task, "test.jsonl"
          )
          return (
              train_dataset,
              val_dataset,
              test_dataset,
              codebase_dataset,
          )
      elif task == "cosqa":
          train_dataset = cosqa_search_train(
              data_dir, task, "cosqa-retrieval-train-19604.json"
          )
          val_dataset = cosqa_search_eval_text(
              data_dir, task, "cosqa-retrieval-dev-500.json"
          )
          test_dataset = cosqa_search_eval_text(
              data_dir, task, "cosqa-retrieval-test-500.json"
          )
          codebase_dataset = cosqa_search_eval_code(data_dir, task)
          return (
              train_dataset,
              val_dataset,
              test_dataset,
              codebase_dataset,
          )
      else:
          train_dataset = csn_search_train(data_dir, task, "train")
          val_dataset = csn_search_eval_text(data_dir, task, "valid")
          test_dataset = csn_search_eval_text(data_dir, task, "test")
          codebase_dataset = csn_search_eval_code(
              data_dir, task, "codebase.jsonl"
          )
          return (
              train_dataset,
              val_dataset,
              test_dataset,
              codebase_dataset,
          )
  
imports: import json
  import torch
  from torch.utils.data import DataLoader, Dataset
label_desc: Use image node_id 1 for calling a global function with example usage: create_dataset(data_dir, task) and returns: train_dataset,  val_dataset,  test_dataset,  codebase_dataset, train_dataset,  val_dataset,  test_dataset,  codebase_dataset, train_dataset,  val_dataset,  test_dataset,  codebase_dataset