wmt16.py 13.1 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 32 33 34 35

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
import tarfile
import gzip
from collections import defaultdict

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

__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 已提交
59
def __build_dict(tar_file, dict_size, save_path, lang):
Y
ying 已提交
60 61 62 63 64 65 66 67 68 69 70 71 72
    word_dict = defaultdict(int)
    with tarfile.open(tar_file, mode="r") as f:
        for line in f.extractfile("wmt16/train"):
            line_split = line.strip().split("\t")
            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(
73 74 75
                    iter(list(word_dict.items())),
                    key=lambda x: x[1],
                    reverse=True)):
Y
ying 已提交
76 77 78 79
            if idx + 3 == dict_size: break
            fout.write("%s\n" % (word[0]))


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

    word_dict = {}
    with open(dict_path, "r") as fdict:
        for idx, line in enumerate(fdict):
            if reverse:
                word_dict[idx] = line.strip()
            else:
                word_dict[line.strip()] = idx
    return word_dict


Y
ying 已提交
97
def __get_dict_size(src_dict_size, trg_dict_size, src_lang):
Y
ying 已提交
98 99 100
    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 已提交
101
                                        TOTAL_EN_WORDS))
Y
ying 已提交
102 103 104 105 106
    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 已提交
107 108 109
        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 已提交
110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 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

        # 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):
                line_split = line.strip().split("\t")
                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 已提交
176 177 178
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type.  Only support: "
                         "en (for English); de(for Germany).")
Y
ying 已提交
179 180
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
181 182

    return reader_creator(
183 184
        tar_file=paddle.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                                "wmt16.tar.gz"),
Y
ying 已提交
185 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
        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 已提交
224 225 226
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type. "
                         "Only support: en (for English); de(for Germany).")
Y
ying 已提交
227

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

    return reader_creator(
232 233
        tar_file=paddle.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                                "wmt16.tar.gz"),
Y
ying 已提交
234 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
        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 已提交
272 273 274
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type. "
                         "Only support: en (for English); de(for Germany).")
Y
ying 已提交
275 276
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
277 278

    return reader_creator(
279 280
        tar_file=paddle.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                                "wmt16.tar.gz"),
Y
ying 已提交
281 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
        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)

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


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.
    """

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