import random


def sentenceize(current_sentences, max_seq_length, target_seq_length):
    """
    Generate a pair of sentences based on a segmentation strategy
    cloned from official electra model.

    Parameters
    ----------
    current_sentences
    max_seq_length
        Maximum sequence length of the training features
    target_seq_length
        Target sequence length of the training features
    Returns
    -------
    first_segment
        The first sentence of the pretraining sequence
    second_segment
        The second sentence of the pretraining sequence.
        Could be None for diversity of training instances.
    """
    # 10% chance to only produce one segment
    if random.random() < 0.1:
        first_segment_target_length = 100000
    else:
        # The reserved space for [CLS] and [SEP] tokens
        first_segment_target_length = (target_seq_length - 3) // 2
    first_segment = []
    second_segment = []
    for sentence in current_sentences:
        if sentence:
            # the sentence goes to the first segment if (1) the first segment is
            # empty, (2) the sentence doesn't put the first segment over length or
            # (3) 50% of the time when it does put the first segment over length
            if (len(first_segment) == 0 or
                    len(first_segment) + len(sentence) < first_segment_target_length or
                    (len(second_segment) == 0 and
                     len(first_segment) < first_segment_target_length and
                     random.random() < 0.5)):
                first_segment += sentence
            else:
                second_segment += sentence

    # trim to max_length while accounting for not-yet-added [CLS]/[SEP] tokens
    first_segment = first_segment[:max_seq_length - 2]
    second_segment = second_segment[:max(0, max_seq_length -
                                         len(first_segment) - 3)]

    return first_segment, second_segment