Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
utils.py 3.94 KB
Newer Older
Zihao Ye's avatar
[Model]Transformer (#186)  
Zihao Ye committed
1
2
3
4
5
6 
import numpy as np
import torch as th
import os
from dgl.data.utils import *

_urls = {
Jinjing Zhou's avatar
fix s3 link (#1310)  
Jinjing Zhou committed
7
8 
    'wmt': 'https://data.dgl.ai/dataset/wmt14bpe_de_en.zip',
    'scripts': 'https://data.dgl.ai/dataset/transformer_scripts.zip',
Zihao Ye's avatar
[Model]Transformer (#186)  
Zihao Ye committed
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25 
}

def prepare_dataset(dataset_name):
    "download and generate datasets"
    script_dir = os.path.join('scripts')
    if not os.path.exists(script_dir):
        download(_urls['scripts'], path='scripts.zip')
        extract_archive('scripts.zip', 'scripts')

    directory = os.path.join('data', dataset_name)
    if not os.path.exists(directory):
        os.makedirs(directory)
    else:
        return
    if dataset_name == 'multi30k':
        os.system('bash scripts/prepare-multi30k.sh')
    elif dataset_name == 'wmt14':
Zihao Ye's avatar
[Model] fix link & beam search (#394)  
Zihao Ye committed
26 
        download(_urls['wmt'], path='wmt14.zip')
Zihao Ye's avatar
[Model]Transformer (#186)  
Zihao Ye committed
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 
        os.system('bash scripts/prepare-wmt14.sh')
    elif dataset_name == 'copy' or dataset_name == 'tiny_copy':
        train_size = 9000
        valid_size = 1000
        test_size = 1000
        char_list = [chr(i) for i in range(ord('a'), ord('z') + 1)]
        with open(os.path.join(directory, 'train.in'), 'w') as f_in,\
            open(os.path.join(directory, 'train.out'), 'w') as f_out:
            for i, l in zip(range(train_size), np.random.normal(15, 3, train_size).astype(int)):
                l = max(l, 1)
                line = ' '.join(np.random.choice(char_list, l)) + '\n'
                f_in.write(line)
                f_out.write(line)

        with open(os.path.join(directory, 'valid.in'), 'w') as f_in,\
            open(os.path.join(directory, 'valid.out'), 'w') as f_out:
            for i, l in zip(range(valid_size), np.random.normal(15, 3, valid_size).astype(int)):
                l = max(l, 1)
                line = ' '.join(np.random.choice(char_list, l)) + '\n'
                f_in.write(line)
                f_out.write(line)

        with open(os.path.join(directory, 'test.in'), 'w') as f_in,\
            open(os.path.join(directory, 'test.out'), 'w') as f_out:
            for i, l in zip(range(test_size), np.random.normal(15, 3, test_size).astype(int)):
                l = max(l, 1)
                line = ' '.join(np.random.choice(char_list, l)) + '\n'
                f_in.write(line)
                f_out.write(line)

        with open(os.path.join(directory, 'vocab.txt'), 'w') as f:
            for c in char_list:
                f.write(c + '\n')

    elif dataset_name == 'sort' or dataset_name == 'tiny_sort':
        train_size = 9000
        valid_size = 1000
        test_size = 1000
        char_list = [chr(i) for i in range(ord('a'), ord('z') + 1)]
        with open(os.path.join(directory, 'train.in'), 'w') as f_in,\
            open(os.path.join(directory, 'train.out'), 'w') as f_out:
            for i, l in zip(range(train_size), np.random.normal(15, 3, train_size).astype(int)):
                l = max(l, 1)
                seq = np.random.choice(char_list, l)
                f_in.write(' '.join(seq) + '\n')
                f_out.write(' '.join(np.sort(seq)) + '\n')

        with open(os.path.join(directory, 'valid.in'), 'w') as f_in,\
            open(os.path.join(directory, 'valid.out'), 'w') as f_out:
            for i, l in zip(range(valid_size), np.random.normal(15, 3, valid_size).astype(int)):
                l = max(l, 1)
                seq = np.random.choice(char_list, l)
                f_in.write(' '.join(seq) + '\n')
                f_out.write(' '.join(np.sort(seq)) + '\n')

        with open(os.path.join(directory, 'test.in'), 'w') as f_in,\
            open(os.path.join(directory, 'test.out'), 'w') as f_out:
            for i, l in zip(range(test_size), np.random.normal(15, 3, test_size).astype(int)):
                l = max(l, 1)
                seq = np.random.choice(char_list, l)
                f_in.write(' '.join(seq) + '\n')
                f_out.write(' '.join(np.sort(seq)) + '\n')

        with open(os.path.join(directory, 'vocab.txt'), 'w') as f:
            for c in char_list:
                f.write(c + '\n')