wmt16.py 13.3 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
import tarfile
from collections import defaultdict

36
import paddle
M
minqiyang 已提交
37
import paddle.compat as cpt
38
import paddle.utils.deprecated as deprecated
Y
ying 已提交
39

40 41
__all__ = []

42
DATA_URL = ("http://paddlemodels.bj.bcebos.com/wmt/wmt16.tar.gz")
Y
ying 已提交
43 44 45 46 47 48 49 50 51 52
DATA_MD5 = "0c38be43600334966403524a40dcd81e"

TOTAL_EN_WORDS = 11250
TOTAL_DE_WORDS = 19220

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


Y
ying 已提交
53
def __build_dict(tar_file, dict_size, save_path, lang):
Y
ying 已提交
54 55 56
    word_dict = defaultdict(int)
    with tarfile.open(tar_file, mode="r") as f:
        for line in f.extractfile("wmt16/train"):
M
minqiyang 已提交
57 58
            line = cpt.to_text(line)
            line_split = line.strip().split("\t")
Y
ying 已提交
59 60 61 62 63
            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

M
minqiyang 已提交
64 65 66
    with open(save_path, "wb") as fout:
        fout.write(
            cpt.to_bytes("%s\n%s\n%s\n" % (START_MARK, END_MARK, UNK_MARK)))
Y
ying 已提交
67
        for idx, word in enumerate(
68 69 70
                sorted(six.iteritems(word_dict),
                       key=lambda x: x[1],
                       reverse=True)):
Y
ying 已提交
71
            if idx + 3 == dict_size: break
M
minqiyang 已提交
72 73
            fout.write(cpt.to_bytes(word[0]))
            fout.write(cpt.to_bytes('\n'))
Y
ying 已提交
74 75


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

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


Y
ying 已提交
93
def __get_dict_size(src_dict_size, trg_dict_size, src_lang):
94 95 96 97
    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 TOTAL_EN_WORDS))
Y
ying 已提交
98 99 100 101
    return src_dict_size, trg_dict_size


def reader_creator(tar_file, file_name, src_dict_size, trg_dict_size, src_lang):
102

Y
ying 已提交
103
    def reader():
Y
ying 已提交
104 105 106
        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 已提交
107

T
tianshuo78520a 已提交
108
        # the index for start mark, end mark, and unk are the same in source
Y
ying 已提交
109 110 111 112 113 114 115 116 117 118 119
        # 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 已提交
120 121
                line = cpt.to_text(line)
                line_split = line.strip().split("\t")
Y
ying 已提交
122 123 124
                if len(line_split) != 2:
                    continue
                src_words = line_split[src_col].split()
125 126 127
                src_ids = [start_id
                           ] + [src_dict.get(w, unk_id)
                                for w in src_words] + [end_id]
Y
ying 已提交
128 129 130 131 132 133 134 135 136 137 138 139

                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


140 141 142
@deprecated(
    since="2.0.0",
    update_to="paddle.text.datasets.WMT16",
143
    level=1,
144
    reason="Please use new dataset API which supports paddle.io.DataLoader")
Y
ying 已提交
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 177 178
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 已提交
179 180 181
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type.  Only support: "
                         "en (for English); de(for Germany).")
Y
ying 已提交
182 183
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
184

185 186 187 188 189 190
    return reader_creator(tar_file=paddle.dataset.common.download(
        DATA_URL, "wmt16", DATA_MD5, "wmt16.tar.gz"),
                          file_name="wmt16/train",
                          src_dict_size=src_dict_size,
                          trg_dict_size=trg_dict_size,
                          src_lang=src_lang)
Y
ying 已提交
191 192


193 194 195
@deprecated(
    since="2.0.0",
    update_to="paddle.text.datasets.WMT16",
196
    level=1,
197
    reason="Please use new dataset API which supports paddle.io.DataLoader")
Y
ying 已提交
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 225 226 227 228 229 230
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 已提交
231 232 233
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type. "
                         "Only support: en (for English); de(for Germany).")
Y
ying 已提交
234

Y
ying 已提交
235 236
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
237

238 239 240 241 242 243
    return reader_creator(tar_file=paddle.dataset.common.download(
        DATA_URL, "wmt16", DATA_MD5, "wmt16.tar.gz"),
                          file_name="wmt16/test",
                          src_dict_size=src_dict_size,
                          trg_dict_size=trg_dict_size,
                          src_lang=src_lang)
Y
ying 已提交
244 245


246 247 248
@deprecated(
    since="2.0.0",
    update_to="paddle.text.datasets.WMT16",
249
    level=1,
250
    reason="Please use new dataset API which supports paddle.io.DataLoader")
Y
ying 已提交
251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282
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 已提交
283 284 285
    if src_lang not in ["en", "de"]:
        raise ValueError("An error language type. "
                         "Only support: en (for English); de(for Germany).")
Y
ying 已提交
286 287
    src_dict_size, trg_dict_size = __get_dict_size(src_dict_size, trg_dict_size,
                                                   src_lang)
Y
ying 已提交
288

289 290 291 292 293 294
    return reader_creator(tar_file=paddle.dataset.common.download(
        DATA_URL, "wmt16", DATA_MD5, "wmt16.tar.gz"),
                          file_name="wmt16/val",
                          src_dict_size=src_dict_size,
                          trg_dict_size=trg_dict_size,
                          src_lang=src_lang)
Y
ying 已提交
295 296


297 298 299
@deprecated(
    since="2.0.0",
    update_to="paddle.text.datasets.WMT16",
300
    level=1,
301
    reason="Please use new dataset API which supports paddle.io.DataLoader")
Y
ying 已提交
302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322
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)

323
    dict_path = os.path.join(paddle.dataset.common.DATA_HOME,
Y
ying 已提交
324
                             "wmt16/%s_%d.dict" % (lang, dict_size))
L
Luo Tao 已提交
325 326 327
    assert os.path.exists(dict_path), "Word dictionary does not exist. "
    "Please invoke paddle.dataset.wmt16.train/test/validation first "
    "to build the dictionary."
328
    tar_file = os.path.join(paddle.dataset.common.DATA_HOME, "wmt16.tar.gz")
Y
ying 已提交
329
    return __load_dict(tar_file, dict_size, lang, reverse)
Y
ying 已提交
330 331


332 333 334
@deprecated(
    since="2.0.0",
    update_to="paddle.text.datasets.WMT16",
335
    level=1,
336
    reason="Please use new dataset API which supports paddle.io.DataLoader")
Y
ying 已提交
337 338 339 340 341
def fetch():
    """download the entire dataset.
    """
    paddle.v4.dataset.common.download(DATA_URL, "wmt16", DATA_MD5,
                                      "wmt16.tar.gz")