wmt14.py 5.9 KB
Newer Older
H
Helin Wang 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14
# Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Q
qijun 已提交
15
WMT14 dataset.
Q
qijun 已提交
16 17
The original WMT14 dataset is too large and a small set of data for set is
provided. This module will download dataset from
Q
qijun 已提交
18
http://paddlepaddle.cdn.bcebos.com/demo/wmt_shrinked_data/wmt14.tgz and
Q
qijun 已提交
19
parse training set and test set into paddle reader creators.
Q
qijun 已提交
20

H
Helin Wang 已提交
21
"""
Q
qiaolongfei 已提交
22
import tarfile
L
Luo Tao 已提交
23
import gzip
Q
qiaolongfei 已提交
24

R
root 已提交
25
import paddle.v2.dataset.common
L
Luo Tao 已提交
26
from paddle.v2.parameters import Parameters
H
Helin Wang 已提交
27

Y
ying 已提交
28 29 30 31 32 33
__all__ = [
    'train',
    'test',
    'get_dict',
    'convert',
]
H
Helin Wang 已提交
34

Y
ying 已提交
35 36
URL_DEV_TEST = ('http://www-lium.univ-lemans.fr/~schwenk/'
                'cslm_joint_paper/data/dev+test.tgz')
H
Helin Wang 已提交
37
MD5_DEV_TEST = '7d7897317ddd8ba0ae5c5fa7248d3ff5'
Y
ying 已提交
38 39 40 41
# this is a small set of data for test. The original data is too large and
# will be add later.
URL_TRAIN = ('http://paddlepaddle.cdn.bcebos.com/demo/'
             'wmt_shrinked_data/wmt14.tgz')
L
Luo Tao 已提交
42
MD5_TRAIN = '0791583d57d5beb693b9414c5b36798c'
43
# BLEU of this trained model is 26.92
L
Luo Tao 已提交
44
URL_MODEL = 'http://paddlepaddle.bj.bcebos.com/demo/wmt_14/wmt14_model.tar.gz'
45
MD5_MODEL = '0cb4a5366189b6acba876491c8724fa3'
Q
qiaolongfei 已提交
46 47 48 49 50 51

START = "<s>"
END = "<e>"
UNK = "<unk>"
UNK_IDX = 2

Q
qiaolongfei 已提交
52 53 54

def __read_to_dict__(tar_file, dict_size):
    def __to_dict__(fd, size):
Q
qiaolongfei 已提交
55
        out_dict = dict()
Q
qiaolongfei 已提交
56 57
        for line_count, line in enumerate(fd):
            if line_count < size:
Q
qiaolongfei 已提交
58 59 60
                out_dict[line.strip()] = line_count
            else:
                break
Q
qiaolongfei 已提交
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
        return out_dict

    with tarfile.open(tar_file, mode='r') as f:
        names = [
            each_item.name for each_item in f
            if each_item.name.endswith("src.dict")
        ]
        assert len(names) == 1
        src_dict = __to_dict__(f.extractfile(names[0]), dict_size)
        names = [
            each_item.name for each_item in f
            if each_item.name.endswith("trg.dict")
        ]
        assert len(names) == 1
        trg_dict = __to_dict__(f.extractfile(names[0]), dict_size)
        return src_dict, trg_dict


def reader_creator(tar_file, file_name, dict_size):
    def reader():
        src_dict, trg_dict = __read_to_dict__(tar_file, dict_size)
        with tarfile.open(tar_file, mode='r') as f:
            names = [
                each_item.name for each_item in f
                if each_item.name.endswith(file_name)
H
Helin Wang 已提交
86
            ]
Q
qiaolongfei 已提交
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 112 113 114
            for name in names:
                for line in f.extractfile(name):
                    line_split = line.strip().split('\t')
                    if len(line_split) != 2:
                        continue
                    src_seq = line_split[0]  # one source sequence
                    src_words = src_seq.split()
                    src_ids = [
                        src_dict.get(w, UNK_IDX)
                        for w in [START] + src_words + [END]
                    ]

                    trg_seq = line_split[1]  # one target sequence
                    trg_words = trg_seq.split()
                    trg_ids = [trg_dict.get(w, UNK_IDX) for w in trg_words]

                    # remove sequence whose length > 80 in training mode
                    if len(src_ids) > 80 or len(trg_ids) > 80:
                        continue
                    trg_ids_next = trg_ids + [trg_dict[END]]
                    trg_ids = [trg_dict[START]] + trg_ids

                    yield src_ids, trg_ids, trg_ids_next

    return reader


def train(dict_size):
Q
qijun 已提交
115
    """
Q
qijun 已提交
116
    WMT14 training set creator.
Q
qijun 已提交
117

Q
qijun 已提交
118 119 120
    It returns a reader creator, each sample in the reader is source language
    word ID sequence, target language word ID sequence and next word ID
    sequence.
Q
qijun 已提交
121

Q
qijun 已提交
122
    :return: Training reader creator
Q
qijun 已提交
123 124
    :rtype: callable
    """
Q
qiaolongfei 已提交
125
    return reader_creator(
R
root 已提交
126 127
        paddle.v2.dataset.common.download(URL_TRAIN, 'wmt14', MD5_TRAIN),
        'train/train', dict_size)
Q
qiaolongfei 已提交
128 129 130


def test(dict_size):
Q
qijun 已提交
131 132 133
    """
    WMT14 test set creator.

Q
qijun 已提交
134 135 136
    It returns a reader creator, each sample in the reader is source language
    word ID sequence, target language word ID sequence and next word ID
    sequence.
Q
qijun 已提交
137

Q
qijun 已提交
138
    :return: Test reader creator
Q
qijun 已提交
139 140
    :rtype: callable
    """
Q
qiaolongfei 已提交
141
    return reader_creator(
R
root 已提交
142 143
        paddle.v2.dataset.common.download(URL_TRAIN, 'wmt14', MD5_TRAIN),
        'test/test', dict_size)
Y
Yancey1989 已提交
144 145


L
Luo Tao 已提交
146 147
def gen(dict_size):
    return reader_creator(
R
root 已提交
148 149
        paddle.v2.dataset.common.download(URL_TRAIN, 'wmt14', MD5_TRAIN),
        'gen/gen', dict_size)
L
Luo Tao 已提交
150 151


L
Luo Tao 已提交
152
def model():
R
root 已提交
153
    tar_file = paddle.v2.dataset.common.download(URL_MODEL, 'wmt14', MD5_MODEL)
L
Luo Tao 已提交
154 155 156 157 158
    with gzip.open(tar_file, 'r') as f:
        parameters = Parameters.from_tar(f)
    return parameters


L
Luo Tao 已提交
159 160 161
def get_dict(dict_size, reverse=True):
    # if reverse = False, return dict = {'a':'001', 'b':'002', ...}
    # else reverse = true, return dict = {'001':'a', '002':'b', ...}
R
root 已提交
162
    tar_file = paddle.v2.dataset.common.download(URL_TRAIN, 'wmt14', MD5_TRAIN)
L
Luo Tao 已提交
163
    src_dict, trg_dict = __read_to_dict__(tar_file, dict_size)
L
Luo Tao 已提交
164 165 166 167
    if reverse:
        src_dict = {v: k for k, v in src_dict.items()}
        trg_dict = {v: k for k, v in trg_dict.items()}
    return src_dict, trg_dict
L
Luo Tao 已提交
168 169


170
def fetch():
R
root 已提交
171 172 173 174 175 176 177 178 179
    paddle.v2.dataset.common.download(URL_TRAIN, 'wmt14', MD5_TRAIN)
    paddle.v2.dataset.common.download(URL_MODEL, 'wmt14', MD5_MODEL)


def convert(path):
    """
    Converts dataset to recordio format
    """
    dict_size = 30000
180 181 182
    paddle.v2.dataset.common.convert(path,
                                     train(dict_size), 1000, "wmt14_train")
    paddle.v2.dataset.common.convert(path, test(dict_size), 1000, "wmt14_test")