update reader (PaddlePaddle#134)

FrostML · web-flow · commit 131eff7d48ae · 2021-03-15T10:45:08.000+08:00
diff --git a/examples/machine_translation/transformer/reader.py b/examples/machine_translation/transformer/reader.py
@@ -23,7 +23,6 @@
 import paddle.distributed as dist
 from paddlenlp.data import Pad, Vocab
 from paddlenlp.datasets import load_dataset
-from paddlenlp.datasets import WMT14ende
 from paddlenlp.data.sampler import SamplerHelper
 
 
@@ -125,12 +124,13 @@ def convert_samples(sample):
 
 
 def adapt_vocab_size(args):
-    root = None if args.root == "None" else args.root
-    (src_vocab, trg_vocab) = WMT14ende.get_vocab(root=root)
+    dataset = load_dataset('wmt14ende', splits=('test'))
+    src_vocab = Vocab.load_vocabulary(**dataset.vocab_info["bpe"])
+    trg_vocab = src_vocab
+
     padding_vocab = (
         lambda x: (x + args.pad_factor - 1) // args.pad_factor * args.pad_factor
     )
-
     args.src_vocab_size = padding_vocab(len(src_vocab))
     args.trg_vocab_size = padding_vocab(len(trg_vocab))
 
