Data format for TFTrainer for TFGpt2
Author: kiyoungkim1
Pull request: no

```train_dataset``` in TFtrainer needs ```(features, labels)```, but TFGpt2 does not need labels (document in TFGPT2LMHeadModel).
Do I know the data format for TFTrainer for TFGpt2?
I have tried this code, but does not work.

Thanks.

```
   def gen_train():
        for ex in transformed_ds[datasets.Split.TRAIN]:
            yield (
                {
                    'input_ids': ex["input_ids"]
                },
                {
                    'labels': ex["input_ids"]
                 }
            )

    train_types = (
                {
                    "input_ids": tf.int32
                },
                {
                    "labels": tf.int32
                },
            )

    train_shapes = (
                {
                    "input_ids": tf.TensorShape([None])
                },
                {
                    "labels": tf.TensorShape([None])
                },
            )

    train_ds = tf.data.Dataset.from_generator(gen_train, train_types, train_shapes)


    if train_ds is not None:
        train_ds = train_ds.apply(tf.data.experimental.assert_cardinality(len(ds[datasets.Split.TRAIN])))
```

Comments:
Hello!

At first, can you share your Transformers and TensorFlow version please?

@jplu 
I use Transformers 4.1.1 and Tensorflow 2.4 in GCP VM, but I can change versions. I will train this model with v3-8 and v3-32.
My code and command are shown below. It is largely based on run_clm.py and run_tf_text_classification.py.
Thanks.

```
import logging
import os
from dataclasses import dataclass, field
from typing import Dict, Optional

import datasets
import numpy as np
import tensorflow as tf

from transformers import (
    AutoConfig,
    AutoTokenizer,
    TFAutoModel,

    GPT2Config,
    GPT2Tokenizer,
    BertTokenizer,
    TFGPT2LMHeadModel,
    HfArgumentParser,
    PreTrainedTokenizer,
    TFTrainer,
    TFTrainingArguments,
)

def get_tfds(
    train_file: str,
    tokenizer: PreTrainedTokenizer,
    max_seq_length: Optional[int] = None,
):
    files = {}

    if train_file is not None:
        files[datasets.Split.TRAIN] = [train_file]

    ds = datasets.load_dataset("csv", data_files=files)
    features_name = 'content'
    transformed_ds = {}

    for k in files.keys():
        transformed_ds[k] = ds[k].map(
            lambda example: tokenizer.batch_encode_plus(
                example[features_name],
                truncation=True,
                max_length=max_seq_length
            ),
            batched=True,
        )

    def gen_train():
        for ex in transformed_ds[datasets.Split.TRAIN]:
            yield (
                {
                    'input_ids': ex["input_ids"]
                },
                {
                    'labels': ex["input_ids"]
                 }
            )

    train_types = (
                {
                    "input_ids": tf.int32
                },
                {
                    "labels": tf.int32
                },
            )

    train_shapes = (
                {
                    "input_ids": tf.TensorShape([None])
                },
                {
                    "labels": tf.TensorShape([None])
                },
            )

    train_ds = tf.data.Dataset.from_generator(gen_train, train_types, train_shapes)


    if train_ds is not None:
        train_ds = train_ds.apply(tf.data.experimental.assert_cardinality(len(ds[datasets.Split.TRAIN])))

    return train_ds


@dataclass
class DataTrainingArguments:
    """
    Arguments pertaining to what data we are going to input our model for training and eval.
    Using `HfArgumentParser` we can turn this class
    into argparse arguments to be able to specify them on
    the command line.
    """

    train_file: str = field(default=None, metadata={"help": "The path of the training file"})
    dev_file: Optional[str] = field(default=None, metadata={"help": "The path of the development file"})
    test_file: Optional[str] = field(default=None, metadata={"help": "The path of the test file"})
    max_seq_length: int = field(
        default=128,
        metadata={
            "help": "The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded."
        },
    )
    overwrite_cache: bool = field(
        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
    )


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
    """

    model_name_or_path: str = field(
        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
    )
    config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
    )
    tokenizer_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
    )
    use_fast: bool = field(default=False, metadata={"help": "Set this flag to use fast tokenization."})
    # If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
    # or just modify its tokenizer_config.json.
    cache_dir: Optional[str] = field(
        default=None,
        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
    )



# See all possible arguments in src/transformers/training_args.py
# or by passing the --help flag to this script.
# We now keep distinct sets of args, for a cleaner separation of concerns.
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TFTrainingArguments))
model_args, data_args, training_args = parser.parse_args_into_dataclasses()

if (
    os.path.exists(training_args.output_dir)
    and os.listdir(training_args.output_dir)
    and training_args.do_train
    and not training_args.overwrite_output_dir
):
    raise ValueError(
        f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
    )

# tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
    cache_dir=model_args.cache_dir,
)
#tokenizer.add_special_tokens({'pad_token': '[PAD]'})

# data
train_dataset = get_tfds(
    train_file=data_args.train_file,
    tokenizer=tokenizer,
    max_seq_length=data_args.max_seq_length,
)

# config
config_kwargs = {
    "cache_dir": model_args.cache_dir,
    #"use_auth_token": True if model_args.use_auth_token else None,
}

config = GPT2Config.from_pretrained(model_args.model_name_or_path, **config_kwargs)

# model
with training_args.strategy.scope():
    model = TFGPT2LMHeadModel.from_pretrained(
        model_args.model_name_or_path,
        from_pt=bool(".bin" in model_args.model_name_or_path),
        config=config,
        cache_dir=model_args.cache_dir,
    )
# model.resize_token_embeddings(len(tokenizer))

# Initialize our Trainer
trainer = TFTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset
)

# Training
if training_args.do_train:
    model_path = (
        model_args.model_name_or_path
        if (model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path))
        else None
    )
    train_result = trainer.train()
    trainer.save_model(model_path)  # Saves the tokenizer too for easy upload
    #tokenizer.save_pretrained(training_args.output_dir)

    output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
    if trainer.is_world_process_zero():
        with open(output_train_file, "w") as writer:
            logger.info("***** Train results *****")
            for key, value in sorted(train_result.metrics.items()):
                logger.info(f"  {key} = {value}")
                writer.write(f"{key} = {value}\n")

        # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
        trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
```

```
python3 run_tftrainer.py \
       --train_file datasets.csv \
       --model_name_or_path gpt2 \
       --do_train \
       --output_dir model \
       --num_train_epochs 4 \
       --per_device_train_batch_size 4 \
       --logging_steps 10 \
       --save_steps 10  \
      --overwrite_output_dir \
      --max_seq_length 128
```



Ok, thanks a lot for sharing this!

If you are using the 4.1.1 release of Transformers, the `TFGPT2LMHeadModel` has a `labels` argument so the problem might come from elsewhere. The other thing to know is that it is currently not possible to train an LM from scratch with TF until the next release (coming very soon), only fine tuning is possible for now.

What is the error you get exactly?

I see.
Error message is shown below. I have tuned many things, but the loss shows nan at best.
Will it take more than a week for the next release? Then I will use pytorch/xla.


```
All the layers of TFGPT2LMHeadModel were initialized from the model checkpoint at gpt2.
If your task is similar to the task the model of the checkpoint was trained on, you can already use TFGPT2LMHeadModel for predictions without further training.
Traceback (most recent call last):
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/eager/context.py", line 2102, in execution_mode
    yield
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 758, in _next_internal
    output_shapes=self._flat_output_shapes)
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 2610, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 6843, in raise_from_not_ok_status
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [65] and element 1 had shape [31].
         [[{{node MultiDeviceIteratorGetNextFromShard}}]]
         [[RemoteCall]] [Op:IteratorGetNext]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "run_tftrainer.py", line 198, in <module>
    train_result = trainer.train()
  File "/home/kykim/.local/lib/python3.6/site-packages/transformers/trainer_tf.py", line 548, in train
    for step, batch in enumerate(train_ds):
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/distribute/input_lib.py", line 649, in __next__
    return self.get_next()
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/distribute/input_lib.py", line 694, in get_next
    self._iterators[i].get_next_as_list_static_shapes(new_name))
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/distribute/input_lib.py", line 1474, in get_next_as_list_static_shapes
    return self._iterator.get_next()
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/data/ops/multi_device_iterator_ops.py", line 581, in get_next
    result.append(self._device_iterators[i].get_next())
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 825, in get_next
    return self._next_internal()
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 764, in _next_internal
    return structure.from_compatible_tensor_list(self._element_spec, ret)
  File "/usr/lib/python3.6/contextlib.py", line 99, in __exit__
    self.gen.throw(type, value, traceback)
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/eager/context.py", line 2105, in execution_mode
    executor_new.wait()
  File "/home/kykim/.local/lib/python3.6/site-packages/tensorflow/python/eager/executor.py", line 67, in wait
    pywrap_tfe.TFE_ExecutorWaitForAllPendingNodes(self._handle)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [65] and element 1 had shape [31].
         [[{{node MultiDeviceIteratorGetNextFromShard}}]]
         [[RemoteCall]]
2021-01-12 19:57:25.672632: W tensorflow/core/kernels/data/generator_dataset_op.cc:103] Error occurred when finalizing GeneratorDataset iterator: Failed precondition: Python interpreter state is not initialized. The process may be terminated.
         [[{{node PyFunc}}]]
```

I know that `tf.data.Dataset.from_generator` has some issues on TPUs, can you rewrite your data processing function to use `tf.data.Dataset.from_tensor_slices` instead?

Also, to be sure that your data are properly formatted you can add an assert to checks this. In order to know if the problems comes from there or not.