wmt16.py 13.2 KB
Newer Older
Y
ying 已提交
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.
"""
Y
ying 已提交
15 16
ACL2016 Multimodal Machine Translation. Please see this website for more
details: http://www.statmt.org/wmt16/multimodal-task.html#task1
Y
ying 已提交
17 18 19 20 21 22 23 24 25 26 27 28 29 30 31

If you use the dataset created for your task, please cite the following paper:
Multi30K: Multilingual English-German Image Descriptions.

@article{elliott-EtAl:2016:VL16,
 author    = {{Elliott}, D. and {Frank}, S. and {Sima"an}, K. and {Specia}, L.},
 title     = {Multi30K: Multilingual English-German Image Descriptions},
 booktitle = {Proceedings of the 6th Workshop on Vision and Language},
 year      = {2016},
 pages     = {70--74},
 year      = 2016
}
"""

import os
M
minqiyang 已提交
32
import six
Y
ying 已提交
33 34 35 36
import tarfile
import gzip
from collections import defaultdict

37
import paddle.dataset.common
M
minqiyang 已提交
38
import paddle.compat as cpt
Y
ying 已提交
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60

__all__ = [
    "train",
    "test",
    "validation",
    "convert",
    "fetch",
    "get_dict",
]

DATA_URL = ("http://cloud.dlnel.org/filepub/"
            "?uuid=46a0808e-ddd8-427c-bacd-0dbc6d045fed")
DATA_MD5 = "0c38be43600334966403524a40dcd81e"

TOTAL_EN_WORDS = 11250
TOTAL_DE_WORDS = 19220

START_MARK = "<s>"
END_MARK = "<e>"
UNK_MARK = "<unk>"


Y
ying 已提交
61
def __build_dict(tar_file, dict_size, save_path, lang):
Y
ying 已提交
62 63 64
    word_dict = defaultdict(int)
    with tarfile.open(tar_file, mode="r") as f:
        for line in f.extractfile("wmt16/train"):
M
minqiyang 已提交
65
            line_split = line.strip().split(six.b("\t"))
Y
ying 已提交
66 67 68 69 70 71 72 73 74
            if len(line_split) != 2: continue
            sen = line_split[0] if lang == "en" else line_split[1]
            for w in sen.split():
                word_dict[w] += 1

    with open(save_path, "w") as fout:
        fout.write("%s\n%s\n%s\n" % (START_MARK, END_MARK, UNK_MARK))
        for idx, word in enumerate(
                sorted(
M
minqiyang 已提交
75
                    six.iteritems(word_dict), key=lambda x: x[1],
76
                    reverse=True)):
Y
ying 已提交
77 78 79 80
            if idx + 3 == dict_size: break
            fout.write("%s\n" % (word[0]))


Y
ying 已提交
81
def __load_dict(tar_file, dict_size, lang, reverse=False):
82
    dict_path = os.path.join(paddle.dataset.common.DATA_HOME,
Y
ying 已提交
83 84
                             "wmt16/%s_%d.dict" % (lang, dict_size))
    if not os.path.exists(dict_path) or (
85
            len(open(dict_path, "rb").readlines()) != dict_size):
Y
ying 已提交
86
        __build_dict(tar_file, dict_size, dict_path, lang)
Y
ying 已提交
87 88

    word_dict = {}
89
    with open(dict_path, "rb") as fdict:
Y
ying 已提交
90 91
        for idx, line in enumerate(fdict):
            if reverse:
M
minqiyang 已提交
92
                word_dict[idx] = cpt.to_text(line.strip())
Y
ying 已提交
93
            else:
M
minqiyang 已提交
94
                word_dict[cpt.to_text(line.strip())] = idx
Y
ying 已提交
95 96 97
    return word_dict


Y
ying 已提交
98
def __get_dict_size(src_dict_size, trg_dict_size, src_lang):
Y
ying 已提交
99 100 101
    src_dict_size = min(src_dict_size, (TOTAL_EN_WORDS if src_lang == "en" else
                                        TOTAL_DE_WORDS))
    trg_dict_size = min(trg_dict_size, (TOTAL_DE_WORDS if src_lang == "en" else
W
Wojciech Uss 已提交
102
                                        TOTAL_EN_WORDS))
Y
ying 已提交
103 104 105 106 107
    return src_dict_size, trg_dict_size


def reader_creator(tar_file, file_name, src_dict_size, trg_dict_size, src_lang):
    def reader():
Y
ying 已提交
108 109 110
        src_dict = __load_dict(tar_file, src_dict_size, src_lang)
        trg_dict = __load_dict(tar_file, trg_dict_size,
                               ("de" if src_lang == "en" else "en"))
Y
ying 已提交
111 112 113 114 115 116 117 118 119 120 121 122 123

        # the indice for start mark, end mark, and unk are the same in source
        # language and target language. Here uses the source language
        # dictionary to determine their indices.
        start_id = src_dict[START_MARK]
        end_id = src_dict[END_MARK]
        unk_id = src_dict[UNK_MARK]

        src_col = 0 if src_lang == "en" else 1
        trg_col = 1 - src_col

        with tarfile.open(tar_file, mode="r") as f:
            for line in f.extractfile(file_name):
M
minqiyang 已提交
124
                line_split = line.strip().split(six.b("\t"))
Y
ying 已提交
125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
                if len(line_split) != 2:
                    continue
                src_words = line_split[src_col].split()
                src_ids = [start_id] + [
                    src_dict.get(w, unk_id) for w in src_words
                ] + [end_id]

                trg_words = line_split[trg_col].split()
                trg_ids = [trg_dict.get(w, unk_id) for w in trg_words]

                trg_ids_next = trg_ids + [end_id]
                trg_ids = [start_id] + trg_ids

                yield src_ids, trg_ids, trg_ids_next

    return reader


def train(src_dict_size, trg_dict_size, src_lang="en"):
    """
    WMT16 train set reader.

    This function returns the reader for train data. Each sample the reader
    returns is made up of three fields: the source language word index sequence,
    target language word index sequence and next word index sequence.


    NOTE:
    The original like for training data is:
    http://www.quest.dcs.shef.ac.uk/wmt16_files_mmt/training.tar.gz

    paddle.dataset.wmt16 provides a tokenized version of the original dataset by
    using moses's tokenization script:
    https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/tokenizer.perl

    Args:
        src_dict_size(int): Size of the source language dictionary. Three
                            special tokens will be added into the dictionary:
                            <s> for start mark, <e> for end mark, and <unk> for
                            unknown word.
        trg_dict_size(int): Size of the target language dictionary. Three
                            special tokens will be added into the dictionary:
                            <s> for start mark, <e> for end mark, and <unk> for
                            unknown word.
        src_lang(string): A string indicating which language is the source
                          language. Available options are: "en" for English
                          and "de" for Germany.

    Returns:
        callable: The train reader.
    """

Y
fix ci.  
ying 已提交
177 178 179
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type.  Only support: "
                         "en (for English); de(for Germany).")
Y
ying 已提交
180 181
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
182 183

    return reader_creator(
184 185
        tar_file=paddle.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                                "wmt16.tar.gz"),
Y
ying 已提交
186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224
        file_name="wmt16/train",
        src_dict_size=src_dict_size,
        trg_dict_size=trg_dict_size,
        src_lang=src_lang)


def test(src_dict_size, trg_dict_size, src_lang="en"):
    """
    WMT16 test set reader.

    This function returns the reader for test data. Each sample the reader
    returns is made up of three fields: the source language word index sequence,
    target language word index sequence and next word index sequence.

    NOTE:
    The original like for test data is:
    http://www.quest.dcs.shef.ac.uk/wmt16_files_mmt/mmt16_task1_test.tar.gz

    paddle.dataset.wmt16 provides a tokenized version of the original dataset by
    using moses's tokenization script:
    https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/tokenizer.perl

    Args:
        src_dict_size(int): Size of the source language dictionary. Three
                            special tokens will be added into the dictionary:
                            <s> for start mark, <e> for end mark, and <unk> for
                            unknown word.
        trg_dict_size(int): Size of the target language dictionary. Three
                            special tokens will be added into the dictionary:
                            <s> for start mark, <e> for end mark, and <unk> for
                            unknown word.
        src_lang(string): A string indicating which language is the source
                          language. Available options are: "en" for English
                          and "de" for Germany.

    Returns:
        callable: The test reader.
    """

Y
fix ci.  
ying 已提交
225 226 227
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type. "
                         "Only support: en (for English); de(for Germany).")
Y
ying 已提交
228

Y
ying 已提交
229 230
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
231 232

    return reader_creator(
233 234
        tar_file=paddle.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                                "wmt16.tar.gz"),
Y
ying 已提交
235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272
        file_name="wmt16/test",
        src_dict_size=src_dict_size,
        trg_dict_size=trg_dict_size,
        src_lang=src_lang)


def validation(src_dict_size, trg_dict_size, src_lang="en"):
    """
    WMT16 validation set reader.

    This function returns the reader for validation data. Each sample the reader
    returns is made up of three fields: the source language word index sequence,
    target language word index sequence and next word index sequence.

    NOTE:
    The original like for validation data is:
    http://www.quest.dcs.shef.ac.uk/wmt16_files_mmt/validation.tar.gz

    paddle.dataset.wmt16 provides a tokenized version of the original dataset by
    using moses's tokenization script:
    https://github.com/moses-smt/mosesdecoder/blob/master/scripts/tokenizer/tokenizer.perl

    Args:
        src_dict_size(int): Size of the source language dictionary. Three
                            special tokens will be added into the dictionary:
                            <s> for start mark, <e> for end mark, and <unk> for
                            unknown word.
        trg_dict_size(int): Size of the target language dictionary. Three
                            special tokens will be added into the dictionary:
                            <s> for start mark, <e> for end mark, and <unk> for
                            unknown word.
        src_lang(string): A string indicating which language is the source
                          language. Available options are: "en" for English
                          and "de" for Germany.

    Returns:
        callable: The validation reader.
    """
Y
fix ci.  
ying 已提交
273 274 275
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type. "
                         "Only support: en (for English); de(for Germany).")
Y
ying 已提交
276 277
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
278 279

    return reader_creator(
280 281
        tar_file=paddle.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                                "wmt16.tar.gz"),
Y
ying 已提交
282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308
        file_name="wmt16/val",
        src_dict_size=src_dict_size,
        trg_dict_size=trg_dict_size,
        src_lang=src_lang)


def get_dict(lang, dict_size, reverse=False):
    """
    return the word dictionary for the specified language.

    Args:
        lang(string): A string indicating which language is the source
                      language. Available options are: "en" for English
                      and "de" for Germany.
        dict_size(int): Size of the specified language dictionary.
        reverse(bool): If reverse is set to False, the returned python
                       dictionary will use word as key and use index as value.
                       If reverse is set to True, the returned python
                       dictionary will use index as key and word as value.

    Returns:
        dict: The word dictionary for the specific language.
    """

    if lang == "en": dict_size = min(dict_size, TOTAL_EN_WORDS)
    else: dict_size = min(dict_size, TOTAL_DE_WORDS)

309
    dict_path = os.path.join(paddle.dataset.common.DATA_HOME,
Y
ying 已提交
310
                             "wmt16/%s_%d.dict" % (lang, dict_size))
L
Luo Tao 已提交
311 312 313
    assert os.path.exists(dict_path), "Word dictionary does not exist. "
    "Please invoke paddle.dataset.wmt16.train/test/validation first "
    "to build the dictionary."
314
    tar_file = os.path.join(paddle.dataset.common.DATA_HOME, "wmt16.tar.gz")
Y
ying 已提交
315
    return __load_dict(tar_file, dict_size, lang, reverse)
Y
ying 已提交
316 317 318 319 320 321 322 323 324 325 326 327 328


def fetch():
    """download the entire dataset.
    """
    paddle.v4.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                      "wmt16.tar.gz")


def convert(path, src_dict_size, trg_dict_size, src_lang):
    """Converts dataset to recordio format.
    """

329
    paddle.dataset.common.convert(
Y
ying 已提交
330 331 332 333 334 335 336
        path,
        train(
            src_dict_size=src_dict_size,
            trg_dict_size=trg_dict_size,
            src_lang=src_lang),
        1000,
        "wmt16_train")
337
    paddle.dataset.common.convert(
Y
ying 已提交
338 339 340 341 342 343 344
        path,
        test(
            src_dict_size=src_dict_size,
            trg_dict_size=trg_dict_size,
            src_lang=src_lang),
        1000,
        "wmt16_test")
345
    paddle.dataset.common.convert(
Y
ying 已提交
346 347 348 349 350 351 352
        path,
        validation(
            src_dict_size=src_dict_size,
            trg_dict_size=trg_dict_size,
            src_lang=src_lang),
        1000,
        "wmt16_validation")