forked from PaddlePaddle/PaddleNLP
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdata.py
123 lines (95 loc) · 4.24 KB
/
data.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
# Copyright (c) 2020 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.
import io
import os
from functools import partial
import numpy as np
import paddle
from paddlenlp.data import Vocab, Pad
from paddlenlp.data import SamplerHelper
from paddlenlp.datasets import load_dataset
def create_train_loader(args):
batch_size = args.batch_size
max_len = args.max_len
train_ds, dev_ds = load_dataset('iwslt15', splits=('train', 'dev'))
src_vocab = Vocab.load_vocabulary(**train_ds.vocab_info['en'])
tgt_vocab = Vocab.load_vocabulary(**train_ds.vocab_info['vi'])
bos_id = src_vocab[src_vocab.bos_token]
eos_id = src_vocab[src_vocab.eos_token]
pad_id = eos_id
def convert_example(example):
source = example['en'].split()[:max_len]
target = example['vi'].split()[:max_len]
source = src_vocab.to_indices(source)
target = tgt_vocab.to_indices(target)
return source, target
key = (lambda x, data_source: len(data_source[x][0]))
# Truncate and convert example to ids
train_ds = train_ds.map(convert_example, lazy=False)
dev_ds = dev_ds.map(convert_example, lazy=False)
train_batch_sampler = SamplerHelper(train_ds).shuffle().sort(
key=key, buffer_size=batch_size * 20).batch(batch_size=batch_size)
dev_batch_sampler = SamplerHelper(dev_ds).sort(
key=key, buffer_size=batch_size * 20).batch(batch_size=batch_size)
train_loader = paddle.io.DataLoader(
train_ds,
batch_sampler=train_batch_sampler,
collate_fn=partial(
prepare_train_input, bos_id=bos_id, eos_id=eos_id, pad_id=pad_id))
dev_loader = paddle.io.DataLoader(
dev_ds,
batch_sampler=dev_batch_sampler,
collate_fn=partial(
prepare_train_input, bos_id=bos_id, eos_id=eos_id, pad_id=pad_id))
return train_loader, dev_loader, len(src_vocab), len(tgt_vocab), pad_id
def create_infer_loader(args):
batch_size = args.batch_size
max_len = args.max_len
test_ds = load_dataset('iwslt15', splits='test')
src_vocab = Vocab.load_vocabulary(**test_ds.vocab_info['en'])
tgt_vocab = Vocab.load_vocabulary(**test_ds.vocab_info['vi'])
bos_id = src_vocab[src_vocab.bos_token]
eos_id = src_vocab[src_vocab.eos_token]
pad_id = eos_id
def convert_example(example):
source = example['en'].split()
target = example['vi'].split()
source = src_vocab.to_indices(source)
target = tgt_vocab.to_indices(target)
return source, target
test_ds.map(convert_example)
test_batch_sampler = SamplerHelper(test_ds).batch(batch_size=batch_size)
test_loader = paddle.io.DataLoader(
test_ds,
batch_sampler=test_batch_sampler,
collate_fn=partial(
prepare_infer_input, bos_id=bos_id, eos_id=eos_id, pad_id=pad_id))
return test_loader, len(src_vocab), len(tgt_vocab), bos_id, eos_id
def prepare_infer_input(insts, bos_id, eos_id, pad_id):
insts = [([bos_id] + inst[0] + [eos_id], [bos_id] + inst[1] + [eos_id])
for inst in insts]
src, src_length = Pad(pad_val=pad_id, ret_length=True)(
[inst[0] for inst in insts])
return src, src_length
def prepare_train_input(insts, bos_id, eos_id, pad_id):
# Add eos token id and bos token id.
insts = [([bos_id] + inst[0] + [eos_id], [bos_id] + inst[1] + [eos_id])
for inst in insts]
# Pad sequence using eos id.
src, src_length = Pad(pad_val=pad_id, ret_length=True)(
[inst[0] for inst in insts])
tgt, tgt_length = Pad(pad_val=pad_id, ret_length=True, dtype="int64")(
[inst[1] for inst in insts])
tgt_mask = (tgt[:, :-1] != pad_id).astype("float32")
return src, src_length, tgt[:, :-1], tgt[:, 1:, np.newaxis], tgt_mask