def get_segments(tokens, max_seq_len):
    """Segments: 0 for the first sequence, 1 for the second"""
    if len(tokens) > max_seq_len:
        raise IndexError('Token length more than max seq length!')
    segments = []
    current_segment_id = 0
    for token in tokens:
        segments.append(current_segment_id)
        if token == '[SEP]':
            current_segment_id = 1
    return segments + [0] * (max_seq_len - len(tokens))
