"""Turning the data from the amazon dataset into something compatible with BERT
"""
def preprocess_data(data, max_num = 100000):
    data_tokens = []
    data_positional = []
    data_targets = []

    #unpacking data
    for i, elem in enumerate(data):

        #tokenizing the title and content
        sentence1 = elem['title']
        sentence2 = elem['content']
        tokens = tokenizer([sentence1, sentence2])
        sentence1_tokens = tokens['input_ids'][0]
        sentence2_tokens = tokens['input_ids'][1]

        # Trimming down tokens
        if len(sentence1_tokens) + len(sentence2_tokens) > max_input_length:
            sentence1_tokens = [101] + sentence1_tokens[-int(max_input_length / 2) + 1:]
            sentence2_tokens = sentence2_tokens[:int(max_input_length / 2) - 1] + [102]

        # Creating sentence tokens
        sentence_tokens = [0] * len(sentence1_tokens) + [1] * len(sentence2_tokens)

        # Combining and padding
        pad_num = max_input_length - (len(sentence1_tokens) + len(sentence2_tokens))
        sequence_tokens = sentence1_tokens + sentence2_tokens + [0] * pad_num
        sentence_location_tokens = sentence_tokens + [1] * pad_num

        data_tokens.append(sequence_tokens)
        data_positional.append(sentence_location_tokens)
        data_targets.append(elem['label'])

        if i > max_num: break

    return torch.tensor(data_positional), torch.tensor(data_tokens), torch.tensor(data_targets)

#processing data into modeling data
train_pos, train_tok, train_targ = preprocess_data(fine_tune_ds['train'])
test_pos, test_tok, test_targ = preprocess_data(fine_tune_ds['test'])

#moving training to device
train_pos = train_pos.to(device)
train_tok = train_tok.to(device)
train_targ = train_targ.to(device)

#moving testing to device
test_pos = test_pos.to(device)
test_tok = test_tok.to(device)
test_targ = test_targ.to(device)