Trainer: To keep unused columns for `compute_metrics` about transformers HOT 4 OPEN

from typing import Dict, List
from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForCausalLM
from transformers import LineByLineTextDataset, DataCollatorForLanguageModeling
from datasets import Dataset

model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-14m")
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-14m")
tokenizer.pad_token = tokenizer.eos_token

data = Dataset.from_dict({'prompt':['A', 'B'], 'completion': ['a', 'b']})

def tokenize(example: Dict[str, str]) -> Dict[str, List[int]]:
  return tokenizer(example['prompt'], f"{example['completion']}.", return_token_type_ids=True)

dataset = data.map(tokenize, remove_columns=['prompt', 'completion'])

args = TrainingArguments(
    output_dir="test",
    report_to='none',
    remove_unused_columns=False,
    max_steps=3,
)

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=args,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
    train_dataset=dataset,
)

trainer.train()