def unify_seq_len(dataset, seq_len, default_id=0):
    """Cut or extend each sequence in dataset to have the same length.
    :param dataset a 2-d array, contains sequences of token ids.
    :param seq_len integer, desired sequence length.
    :param default_id if a sequence is shorter than seq_len, default_id
            will be appended to the end.
    :return a 2-d array, contains sequences of token ids, each sequence
            has length seq_len
    """
    return [sample[:seq_len] if seq_len < len(sample) else
         sample + [default_id for _ in range(0, seq_len - len(sample))]
         for sample in dataset]