code: def tok_text_gen_mlm(
      tokenizer: PreTrainedTokenizer,
      dataset: Dataset,
      args: Union[GENTrainArgs, WPTrainArgs, GENEvalArgs, WPEvalArgs],
      mlm: bool = True,
  ) -> Dataset:
      # todo set to args.max_length
      if args.max_length is None:
          max_input_length = tokenizer.model_max_length
      else:
          max_input_length = args.max_length
  
      def tokenize_function(example):
          texts = example["text"]
          # Add newlines back that were removed from load_dataset().
          texts = [case + "\n" for case in texts[:]]
          return tokenizer(texts)
  
      tokenized_dataset = dataset.map(
          tokenize_function, batched=True, remove_columns=["text"]
      )
  
      def group_texts(examples):
          concatenated_examples = {
              k: sum(examples[k], []) for k in examples.keys()
          }
          total_length = len(
              concatenated_examples[list(examples.keys())[0]]
          )
          output_length = (
              total_length // max_input_length
          ) * max_input_length
  
          # if  total_length is less than the max_input_length length
          # then it causes an error unless the code below is used.
          # this is due to total_length being truncated to 0
          if output_length == 0:
              output_length = total_length
  
          result = {
              k: [
                  t[i : i + max_input_length]
                  for i in range(0, output_length, max_input_length)
              ]
              for k, t in concatenated_examples.items()
          }
  
          if not mlm:
              # Masked language models don't need labels. Text generation models do
              result["labels"] = result["input_ids"].copy()
          return result
  
      tokenized_dataset = tokenized_dataset.map(
          group_texts, batched=True
      )
  
      return tokenized_dataset
  
file_path: happy-transformer/happytransformer/fine_tuning_util.py
imports: from dataclasses import dataclass
  from typing import Union
  from datasets import Dataset
  from transformers import PreTrainedTokenizer, TrainerCallback
  from happytransformer.args import GENTrainArgs, WPTrainArgs, GENEvalArgs, WPEvalArgs
label_desc: Use image node_id 1 for calling a global function with example usage: tok_text_gen_mlm(tokenizer, dataset, args, mlm) and returns: tokenized_dataset, tokenizer, result