imikolov_test.py 2.3 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# 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
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# 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.

15
import paddle.dataset.imikolov
16 17
import unittest

18
WORD_DICT = paddle.dataset.imikolov.build_dict()
19

20 21 22 23 24 25 26 27

class TestMikolov(unittest.TestCase):
    def check_reader(self, reader, n):
        for l in reader():
            self.assertEqual(len(l), n)

    def test_train(self):
        n = 5
28
        self.check_reader(paddle.dataset.imikolov.train(WORD_DICT, n), n)
29

30 31 32 33 34 35 36
        first_line = 'aer banknote berlitz calloway centrust cluett fromstein '\
            'gitano guterman hydro-quebec ipo kia memotec mlx nahb punts '\
            'rake regatta rubens sim snack-food ssangyong swapo wachter'
        first_line = [
            WORD_DICT.get(ch, WORD_DICT['<unk>'])
            for ch in first_line.split(' ')
        ]
37
        for l in paddle.dataset.imikolov.train(
38
                WORD_DICT, n=-1,
39
                data_type=paddle.dataset.imikolov.DataType.SEQ)():
40 41 42 43
            read_line = l[0][1:]
            break
        self.assertEqual(first_line, read_line)

44 45
    def test_test(self):
        n = 5
46
        self.check_reader(paddle.dataset.imikolov.test(WORD_DICT, n), n)
47

48 49 50 51 52 53
        first_line = 'consumers may want to move their telephones a little '\
                'closer to the tv set'
        first_line = [
            WORD_DICT.get(ch, WORD_DICT['<unk>'])
            for ch in first_line.split(' ')
        ]
54
        for l in paddle.dataset.imikolov.test(
55
                WORD_DICT, n=-1,
56
                data_type=paddle.dataset.imikolov.DataType.SEQ)():
57 58 59 60
            read_line = l[0][1:]
            break
        self.assertEqual(first_line, read_line)

61 62 63
    def test_total(self):
        _, idx = zip(*WORD_DICT.items())
        self.assertEqual(sorted(idx)[-1], len(WORD_DICT) - 1)
64 65 66 67


if __name__ == '__main__':
    unittest.main()