import torch


def encode_data(data, tokenizer, punctuation_enc, segment_size):
    """
    Converts words to (BERT) tokens and punctuation to given encoding.
    Note that words can be composed of multiple tokens.
    """

    X = []
    Y = []
    for idx, split_text in enumerate(data):
        if len(split_text) >= segment_size:
            X_tmp = []
            for word, punc in split_text:
                tokens = tokenizer.tokenize(word)
                x = tokenizer.convert_tokens_to_ids(tokens)
                y = [punctuation_enc[punc]]
                if len(x) > 0:
                    if len(x) > 1:
                        y = (len(x)-1)*[0] + y
                    X_tmp += x
                    Y += y
            X.append(X_tmp)

    return X, torch.tensor(Y)