ValueError: The batch received was empty

[cyw3]

I run the demo in README:

# -*- encoding: utf-8 -*-

from codetf.trainer.causal_lm_trainer import CausalLMTrainer
from codetf.data_utility.codexglue_dataset import CodeXGLUEDataset
from codetf.models import load_model_pipeline
from codetf.performance.evaluation_metric import EvaluationMetric


model_class = load_model_pipeline(
    model_name="causal-lm",
    # model_name="codet5",
    task="pretrained",
    # model_type="starcoder-15.5B",
    model_type="codegen-350M-mono",
    # model_type="base-multi-sum",
    is_eval=False,
    load_in_8bit=False,
    weight_sharding=False,
)


dataloader = CodeXGLUEDataset(tokenizer=model_class.get_tokenizer())
train_dataset, test_dataset, val_dataset = dataloader.load(subset="text-to-code")

# peft can be in ["lora", "prefixtuning"]
trainer = CausalLMTrainer(
    train_dataset=train_dataset,
    validation_dataset=val_dataset,
    peft=None,
    pretrained_model_or_path=model_class.get_model(),
    tokenizer=model_class.get_tokenizer(),
)
trainer.train()


evaluator = EvaluationMetric(metric="bleu", tokenizer=model_class.tokenizer)
# trainer.evaluate(test_dataset=test_dataset)

However, I get a error:

╭─────────────────────────────── Traceback (most recent call last) ────────────────────────────────╮
│ /data/test//test_finetune.py:36 in <module>                        │
│                                                                                                  │
│   33 │   pretrained_model_or_path=model_class.get_model(),                                       │
│   34 │   tokenizer=model_class.get_tokenizer(),                                                  │
│   35 )                                                                                           │
│ ❱ 36 trainer.train()                                                                             │
│   37                                                                                             │
│   38                                                                                             │
│   39 evaluator = EvaluationMetric(metric="bleu", tokenizer=model_class.tokenizer)                │
│                                                                                                  │
│ /root/.pyenv/versions/3.10.0/lib/python3.10/site-packages/codetf/trainer/base_trainer.py:54 in train                          │
│                                                                                                  │
│    51 │   │   )                                                                                  │
│    52 │                                                                                          │
│    53 │   def train(self):                                                                       │
│ ❱  54 │   │   self.trainer.train()                                                               │
│    55 │                                                                                          │
│    56 │   def evaluate(self, dataset=None):                                                      │
│    57 │   │   self.trainer.evaluate(dataset)                                                     │
│                                                                                                  │
│ /root/.pyenv/versions/3.10.0/lib/python3.10/site-packages/transformers/trainer.py:1664 in train  │
│                                                                                                  │
│   1661 │   │   inner_training_loop = find_executable_batch_size(                                 │
│   1662 │   │   │   self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size  │
│   1663 │   │   )                                                                                 │
│ ❱ 1664 │   │   return inner_training_loop(                                                       │
│   1665 │   │   │   args=args,                                                                    │
│   1666 │   │   │   resume_from_checkpoint=resume_from_checkpoint,                                │
│   1667 │   │   │   trial=trial,                                                                  │
│                                                                                                  │
│ /root/.pyenv/versions/3.10.0/lib/python3.10/site-packages/accelerate/utils/memory.py:124 in      │
│ decorator                                                                                        │
│                                                                                                  │
│   121 │   │   │   if batch_size == 0:                                                            │
│   122 │   │   │   │   raise RuntimeError("No executable batch size found, reached zero.")        │
│   123 │   │   │   try:                                                                           │
│ ❱ 124 │   │   │   │   return function(batch_size, *args, **kwargs)                               │
│   125 │   │   │   except Exception as e:                                                         │
│   126 │   │   │   │   if should_reduce_batch_size(e):                                            │
│   127 │   │   │   │   │   gc.collect()                                                           │
│                                                                                                  │
│ /root/.pyenv/versions/3.10.0/lib/python3.10/site-packages/transformers/trainer.py:1940 in        │
│ _inner_training_loop                                                                             │
│                                                                                                  │
│   1937 │   │   │   │   │   with model.no_sync():                                                 │
│   1938 │   │   │   │   │   │   tr_loss_step = self.training_step(model, inputs)                  │
│   1939 │   │   │   │   else:                                                                     │
│ ❱ 1940 │   │   │   │   │   tr_loss_step = self.training_step(model, inputs)                      │
│   1941 │   │   │   │                                                                             │
│   1942 │   │   │   │   if (                                                                      │
│   1943 │   │   │   │   │   args.logging_nan_inf_filter                                           │
│                                                                                                  │
│ /root/.pyenv/versions/3.10.0/lib/python3.10/site-packages/transformers/trainer.py:2728 in        │
│ training_step                                                                                    │
│                                                                                                  │
│   2725 │   │   │   `torch.Tensor`: The tensor with training loss on this batch.                  │
│   2726 │   │   """                                                                               │
│   2727 │   │   model.train()                                                                     │
│ ❱ 2728 │   │   inputs = self._prepare_inputs(inputs)                                             │
│   2729 │   │                                                                                     │
│   2730 │   │   if is_sagemaker_mp_enabled():                                                     │
│   2731 │   │   │   loss_mb = smp_forward_backward(model, inputs, self.args.gradient_accumulatio  │
│                                                                                                  │
│ /root/.pyenv/versions/3.10.0/lib/python3.10/site-packages/transformers/trainer.py:2675 in        │
│ _prepare_inputs                                                                                  │
│                                                                                                  │
│   2672 │   │   """                                                                               │
│   2673 │   │   inputs = self._prepare_input(inputs)                                              │
│   2674 │   │   if len(inputs) == 0:                                                              │
│ ❱ 2675 │   │   │   raise ValueError(                                                             │
│   2676 │   │   │   │   "The batch received was empty, your model won't be able to train on it.   │
│   2677 │   │   │   │   f"training dataset contains keys expected by the model: {','.join(self._  │
│   2678 │   │   │   )                                                                             │
╰──────────────────────────────────────────────────────────────────────────────────────────────────╯
ValueError: The batch received was empty, your model won't be able to train on it. Double-check that 
your training dataset contains keys expected by the model: 
input_ids,past_key_values,attention_mask,token_type_ids,position_ids,head_mask,inputs_embeds,labels,us
e_cache,output_attentions,output_hidden_states,return_dict,labels,label,label_ids.
  0%|                                                                          | 0/10 [00:01<?, ?it/s]