def text_gen_pre_process(dataset, tokenizer, all_kwargs):
    """Pre-process data for text generation task.

    The input dataset is expected to have one or more text columns.
    Depending on the strategy, the sequences are either joined together or processed individually.
    """
    import torch
    from datasets import Dataset as HFDataset

    args = validate_config(all_kwargs, TextGenParams, warn_unused_keys=True)

    # set tokenizer padding side
    tokenizer.padding_side = args.padding_side

    if isinstance(args.text_formatting_func, str):
        # load text_formatting_func
        args.text_formatting_func = args.get_user_module_loader().load_object(args.text_formatting_func)
    # get text from dataset
    dataset = dataset.map(
        lambda x: {
            "text": get_text(
                x,
                args.chat_template,
                args.message_col,
                args.text_formatting_func,
                args.text_template,
                args.text_cols,
                args.add_special_tokens,
                tokenizer,
            )
        }
    )
    text_list = [text for text in dataset["text"] if text]  # remove empty strings
    total_examples = len(text_list)  # total number of examples

    tokenized_inputs = {"input_ids": [], "labels": [], "attention_mask": []}
    if "join" in args.strategy:
        joiner_tokens = tokenizer.encode(args.joiner, add_special_tokens=False) if args.joiner else []

        if args.strategy != TextGenStrategy.JOIN_RANDOM:
            # no randomization, just use contiguous blocks of tokens
            if args.strategy == TextGenStrategy.JOIN_SLIDING_WINDOW:
                # we use the stride as both the step between sequences and the context size
                step, context = args.stride, args.max_seq_len - args.stride
            else:
                # JOIN strategy
                # text is split into non-overlapping sequences and there is no context
                step, context = args.max_seq_len, None

            example_idx = 0  # index of the first example in the current batch
            num_samples = 0  # samples processed so far
            overflow = []  # tokens overflowed from the previous batch of examples
            # we will process in batches to make tokenization faster
            # better than joining all text together and tokenizing all at once
            while True:
                if args.max_samples is not None and num_samples >= args.max_samples:
                    # we have reached max_samples
                    break
                if example_idx >= total_examples:
                    # we have reached the end of the text_list
                    break

                examples_to_get = min(args.processing_batch_size, total_examples - example_idx)
                # batch tokenize
                batched_input_ids = tokenizer(
                    text_list[example_idx : example_idx + examples_to_get],
                    add_special_tokens=False,
                    truncation=False,
                )["input_ids"]

                # join all the input_ids together with joiner_tokens
                joined_input_ids = overflow
                for input_ids in batched_input_ids:
                    joined_input_ids += input_ids + joiner_tokens

                end_loc = 0  # position of unused token in joined_input_ids
                # '- args.max_seq_len ' is used to make sure we don't get a sequence that is too short
                for begin_loc in range(0, len(joined_input_ids) - args.max_seq_len, step):
                    # end_loc is the beginning of the next sequence
                    end_loc = begin_loc + args.max_seq_len
                    # get the input sequence
                    input_ids = torch.tensor(joined_input_ids[begin_loc:end_loc])
                    append_text_gen_input_ids(tokenized_inputs, input_ids, torch.ones_like(input_ids), context=context)
                    num_samples += 1
                    if args.max_samples is not None and num_samples >= args.max_samples:
                        # we have reached max_samples
                        break
                # update counters
                example_idx += examples_to_get
                overflow = joined_input_ids[end_loc:]
        else:
            # randomization, sample random blocks of tokens
            rng = Random(args.random_seed)
            # cache to store tokenized examples
            cache = {}
            for _ in range(args.max_samples):
                resamples = 0
                # will try to sample sequences random_retries times before giving up
                while resamples < args.random_retries:
                    # sample a beginning example
                    # randint is inclusive, so we need to subtract 1
                    begin_example_idx = rng.randint(0, total_examples - 1)
                    joined_input_ids = []
                    # loop through the examples until we have enough tokens
                    for i in range(begin_example_idx, total_examples):
                        # get the input_ids
                        if i not in cache:
                            cache[i] = tokenizer(
                                text_list[i],
                                add_special_tokens=False,
                                truncation=False,
                            )["input_ids"]
                        joined_input_ids += cache[i] + joiner_tokens
                        # stop if we have enough tokens
                        if len(joined_input_ids) >= args.max_seq_len:
                            break
                    # add to samples if we have enough tokens
                    if len(joined_input_ids) >= args.max_seq_len:
                        # found a good example
                        input_ids = torch.tensor(joined_input_ids[: args.max_seq_len])
                        append_text_gen_input_ids(tokenized_inputs, input_ids, torch.ones_like(input_ids))
                        break
                    resamples += 1
    else:
        # each line is a sequence
        if args.strategy == TextGenStrategy.LINE_BY_LINE:
            # batched tokenization might be faster so lets tokenize all the text at once
            if not args.max_samples:
                for native_input_ids, native_attention_mask in batch_tokenize_text(text_list, tokenizer, args):
                    append_text_gen_input_ids(
                        tokenized_inputs, torch.tensor(native_input_ids), torch.tensor(native_attention_mask)
                    )
            else:
                example_idx = 0  # index of the first example in the current batch
                num_samples = 0
                while True:
                    if num_samples >= args.max_samples or example_idx >= total_examples:
                        # we have reached max_samples or the end of the text_list
                        break
                    # get as many examples as possible without going over max_samples
                    examples_to_get = min(args.max_samples - num_samples, total_examples - example_idx)
                    # batch tokenize
                    tokenized_texts = batch_tokenize_text(
                        text_list[example_idx : example_idx + examples_to_get],
                        tokenizer,
                        args,
                    )
                    for native_input_ids, native_attention_mask in tokenized_texts:
                        append_text_gen_input_ids(
                            tokenized_inputs, torch.tensor(native_input_ids), torch.tensor(native_attention_mask)
                        )
                        num_samples += 1
                    # update counters
                    example_idx += examples_to_get
        else:
            # randomization, sample random lines
            rng = Random(args.random_seed)
            cache = {}
            for _ in range(args.max_samples):
                resamples = 0
                encodings = None
                while resamples < args.random_retries:
                    # sample a random line
                    # randint is inclusive, so we need to subtract 1
                    i = rng.randint(0, len(text_list) - 1)
                    if i not in cache:
                        encodings = tokenizer(
                            text_list[i],
                            max_length=args.max_seq_len,
                            truncation=True,
                            padding="max_length" if args.pad_to_max_len else False,
                            add_special_tokens=False,
                            return_tensors="pt",
                        )
                        cache[i] = encodings
                    else:
                        encodings = cache[i]
                    if not args.drop_short_sequences or encodings.input_ids.shape[1] >= args.max_seq_len:
                        # found a good sample
                        break
                    resamples += 1
                if not encodings:
                    # could not find a good sample after resampling
                    continue
                append_text_gen_input_ids(tokenized_inputs, encodings.input_ids[0], encodings.attention_mask[0])

    if not args.use_attention_mask:
        # remove attention_mask
        tokenized_inputs.pop("attention_mask")

    # convert to HFDataset
    hf_dataset = HFDataset.from_dict(tokenized_inputs)
    hf_dataset.set_format("torch", output_all_columns=True)

    # return BaseDataset
    return BaseDataset(hf_dataset, "labels", max_samples=args.max_samples)
Pre-process data for text generation task.

    The input dataset is expected to have one or more text columns.
    Depending on the strategy, the sequences are either joined together or processed individually.
    
olive/data/component/text_generation.py