SentiX/data_utils_aspect.py at master · DrJZhou/SentiX · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
import os
import pickle
import numpy as np
import torch
from torch.utils.data import Dataset
from transformers import BertTokenizer
bert_path = "/gruntdata/zhoujie/bert_model"
# bert_path = "D:/Projects/bert_model"


def pad_and_truncate(sequence, maxlen, dtype='int64', padding='post', truncating='post', value=0):
    x = (np.ones(maxlen) * value).astype(dtype)
    if truncating == 'pre':
        trunc = sequence[-maxlen:]
    else:
        trunc = sequence[:maxlen]
    trunc = np.asarray(trunc, dtype=dtype)
    if padding == 'post':
        x[:len(trunc)] = trunc
    else:
        x[-len(trunc):] = trunc
    return x


class Tokenizer4Bert:
    def __init__(self, max_seq_len):
        self.tokenizer = BertTokenizer.from_pretrained(bert_path)
        self.max_seq_len = max_seq_len

    def text_to_sequence(self, text, reverse=False, padding='post', truncating='post'):
        sequence = self.tokenizer.convert_tokens_to_ids(self.tokenizer.tokenize(text))
        if len(sequence) == 0:
            sequence = [0]
        if reverse:
            sequence = sequence[::-1]
        return pad_and_truncate(sequence, self.max_seq_len, padding=padding, truncating=truncating)


class ABSADataset(Dataset):
    def __init__(self, data_name, tokenizer, train_flag="train"):
        fname = './data/aspect-based_sentiment_analysis/{}/{}.raw'.format(data_name, train_flag)
        fin = open(fname, 'r', encoding='utf-8', newline='\n', errors='ignore')
        lines = fin.readlines()
        fin.close()

        all_data = []
        for i in range(0, len(lines), 3):
            text_left, _, text_right = [s.lower().strip() for s in lines[i].partition("$T$")]
            aspect = lines[i + 1].lower().strip()
            polarity = lines[i + 2].strip()
            text_raw_indices = tokenizer.text_to_sequence("[CLS] " + text_left + " " + aspect + " " + text_right + ' [SEP] ')
            aspect_indices = tokenizer.text_to_sequence(aspect)
            text_left_indices = tokenizer.text_to_sequence("[CLS] " + text_left)
            text_left_len = np.sum(text_left_indices != 0)
            aspect_len = np.sum(aspect_indices != 0)
            polarity = int(polarity) + 1
            if text_left_len + aspect_len > tokenizer.max_seq_len:
                continue
            position_indices = np.array([text_left_len, text_left_len+aspect_len])
            text_bert_indices = tokenizer.text_to_sequence('[CLS] ' + text_left + " " + aspect + " " + text_right + ' [SEP] ' + aspect + " [SEP]")
            bert_segments_ids = np.asarray([0] * (np.sum(text_raw_indices != 0) + 2) + [1] * (aspect_len + 1))
            bert_segments_ids = pad_and_truncate(bert_segments_ids, tokenizer.max_seq_len)

            data = {
                'text_raw_indices': text_raw_indices,
                'text_bert_indices': text_bert_indices,
                'bert_segments_ids': bert_segments_ids,
                'position_indices': position_indices,
                'polarity': polarity,
            }
            all_data.append(data)
        self.data = all_data

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return len(self.data)

#
# if __name__ == '__main__':
#     tokenizer = Tokenizer4Bert(100)
#     dataset = ABSADataset('res14', tokenizer, train_flag='train')