def reverse_nn_input(dataset, service_tokens):
    """
    Swaps the last utterance of x with y for each x-y pair in the dataset.
    To handle different length of sequences, everything is filled with pads
    to the length of longest sequence.
    """
    # Swap last utterance of x with y, while padding with start- and eos-tokens
    y_output = np.full(dataset.y.shape, service_tokens.pad_token_id, dtype=dataset.y.dtype)
    for y_output_sample, x_input_sample in zip(y_output, dataset.x[:, -1]):
        # Write start token at the first index
        y_output_sample[0] = service_tokens.start_token_id
        y_output_token_index = 1
        for value in x_input_sample:
            # We should stop at pad tokens in the input sample
            if value == service_tokens.pad_token_id:
                break
            # We should keep last token index with pad, so we can replace it futher with eos-token
            if y_output_token_index == y_output_sample.shape[-1] - 1:
                break
            y_output_sample[y_output_token_index] = value
            y_output_token_index += 1
        # Write eos token right after the last non-pad token in the sample
        y_output_sample[y_output_token_index] = service_tokens.eos_token_id

    # Use utterances from y in x while truncating start- and eos-tokens
    x_output = np.full(dataset.x.shape, service_tokens.pad_token_id, dtype=dataset.x.dtype)
    for x_output_sample, x_input_sample, y_input_sample in zip(x_output, dataset.x[:, :-1], dataset.y):
        # Copy all the context utterances except the last one right to the output
        x_output_sample[:-1] = x_input_sample
        x_output_token_index = 0
        for value in y_input_sample:
            # Skip start- and eos-tokens from the input sample because we don't need them in X
            if value in {service_tokens.start_token_id, service_tokens.eos_token_id}:
                continue
            # Stop if we already reached the end of output sample (in case the input sample is longer than output)
            if x_output_token_index == x_output_sample.shape[-1]:
                break
            # Fill the tokens of the last utterance in dialog context
            x_output_sample[-1, x_output_token_index] = value
            x_output_token_index += 1

    return Dataset(x=x_output, y=y_output, condition_ids=dataset.condition_ids)

    Swaps the last utterance of x with y for each x-y pair in the dataset.
    To handle different length of sequences, everything is filled with pads
    to the length of longest sequence.
    
cakechat/dialog_model/model_utils.py