attention_rnn/data_utils.py at master · Walusus/attention_rnn · 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
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
import os
import json
import torch


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Label encoding dicts
class_to_num = {"AddToPlaylist": 0,
                "BookRestaurant": 1,
                "GetWeather": 2,
                "PlayMusic": 3,
                "RateBook": 4,
                "SearchCreativeWork": 5,
                "SearchScreeningEvent": 6}

num_to_class = {0: "AddToPlaylist",
                1: "BookRestaurant",
                2: "GetWeather",
                3: "PlayMusic",
                4: "RateBook",
                5: "SearchCreativeWork",
                6: "SearchScreeningEvent"}


def string_to_word_list(string):
    word_list = string.split()
    stripped_list = []

    for word in word_list:
        # Cleanup and unify words
        strp_word = word.strip(" .:?!,").rstrip("'s").lower()
        if strp_word.replace('/', '').replace(':', '').replace('-', '').isnumeric():
            strp_word = " num "
        stripped_list.append(strp_word)

    return stripped_list


def load_labeled_data():
    valid_s = []
    train_s = []

    root, _, files = next(os.walk("datasets/"))
    for filename in files:
        with open(root + filename, 'r') as f:
            _, data = json.load(f).popitem()

            # Determine label and data type by file name
            label = class_to_num[filename.strip(".json").split("_")[1]]

            for s in data:
                sentence = ''.join(part['text'] for part in s['data'])

                if filename.startswith("validate"):
                    valid_s.append((sentence, label))
                else:
                    train_s.append((sentence, label))

    return train_s, valid_s


def get_encodings(dataset):
    # Use set to avoid repetition
    vocab = set()
    max_len = 0

    for sentence, _ in dataset:
        word_list = string_to_word_list(sentence)
        max_len = len(word_list) if len(word_list) > max_len else max_len
        vocab.update(word_list)

    word_to_idx = {word: i + 1 for i, word in enumerate(vocab)}
    # Special word to pad sentences
    word_to_idx[" pad "] = 0
    return word_to_idx, max_len


def encode_string(sentence, encodings):
    return [encodings[word] for word in string_to_word_list(sentence)]


def encode_dataset(dataset, encodings):
    encoded_data = []
    for value, label in dataset:
        encoded_data.append((encode_string(value, encodings), label))

    return encoded_data


def sentence_to_tensor(sentence, max_len, label=None):
    # Pad sentence with zeros
    sentence += [0 for i in range(max_len - len(sentence))]
    ts = torch.tensor(sentence, dtype=torch.long, device=device).unsqueeze(0)

    if label is not None:
        tl = torch.tensor(label, dtype=torch.long, device=device).unsqueeze(0)
        return ts, tl
    else:
        return ts


def dataset_to_tensor(dataset, max_len):
    sentence_list = []
    label_list = []
    for value, label in dataset:
        v_t, l_t = sentence_to_tensor(value, max_len, label)
        sentence_list.append(v_t)
        label_list.append(l_t)

    return torch.cat(sentence_list), torch.cat(label_list)