test_tokenization_ctrl.py 2.56 KB
Newer Older
keskarnitish's avatar
keskarnitish committed
1
2
3
4
5
6
7
8
9
10
11
12
13
# coding=utf-8
# Copyright 2018 Salesforce and HuggingFace Inc. team.
# 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.
Aymeric Augustin's avatar
Aymeric Augustin committed
14

keskarnitish's avatar
keskarnitish committed
15

Aymeric Augustin's avatar
Aymeric Augustin committed
16
import json
keskarnitish's avatar
keskarnitish committed
17
import os
18
import unittest
keskarnitish's avatar
keskarnitish committed
19

Aymeric Augustin's avatar
Aymeric Augustin committed
20
from transformers.tokenization_ctrl import VOCAB_FILES_NAMES, CTRLTokenizer
keskarnitish's avatar
keskarnitish committed
21

22
from .test_tokenization_common import TokenizerTesterMixin
keskarnitish's avatar
keskarnitish committed
23

24

25
class CTRLTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
keskarnitish's avatar
keskarnitish committed
26
27

    tokenizer_class = CTRLTokenizer
28
    test_rust_tokenizer = False
keskarnitish's avatar
keskarnitish committed
29
30

    def setUp(self):
Julien Chaumond's avatar
Julien Chaumond committed
31
        super().setUp()
keskarnitish's avatar
keskarnitish committed
32
33

        # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
34
        vocab = ["adapt", "re@@", "a@@", "apt", "c@@", "t", "<unk>"]
keskarnitish's avatar
keskarnitish committed
35
        vocab_tokens = dict(zip(vocab, range(len(vocab))))
36
        merges = ["#version: 0.2", "a p", "ap t</w>", "r e", "a d", "ad apt</w>", ""]
keskarnitish's avatar
keskarnitish committed
37
38
        self.special_tokens_map = {"unk_token": "<unk>"}

39
40
        self.vocab_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
        self.merges_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["merges_file"])
keskarnitish's avatar
keskarnitish committed
41
42
43
44
45
46
47
48
49
        with open(self.vocab_file, "w", encoding="utf-8") as fp:
            fp.write(json.dumps(vocab_tokens) + "\n")
        with open(self.merges_file, "w", encoding="utf-8") as fp:
            fp.write("\n".join(merges))

    def get_tokenizer(self, **kwargs):
        kwargs.update(self.special_tokens_map)
        return CTRLTokenizer.from_pretrained(self.tmpdirname, **kwargs)

50
    def get_input_output_texts(self, tokenizer):
51
52
        input_text = "adapt react readapt apt"
        output_text = "adapt react readapt apt"
keskarnitish's avatar
keskarnitish committed
53
54
55
56
57
        return input_text, output_text

    def test_full_tokenizer(self):
        tokenizer = CTRLTokenizer(self.vocab_file, self.merges_file, **self.special_tokens_map)
        text = "adapt react readapt apt"
58
        bpe_tokens = "adapt re@@ a@@ c@@ t re@@ adapt apt".split()
thomwolf's avatar
thomwolf committed
59
        tokens = tokenizer.tokenize(text)
keskarnitish's avatar
keskarnitish committed
60
61
62
63
64
        self.assertListEqual(tokens, bpe_tokens)

        input_tokens = tokens + [tokenizer.unk_token]

        input_bpe_tokens = [0, 1, 2, 4, 5, 1, 0, 3, 6]
65
        self.assertListEqual(tokenizer.convert_tokens_to_ids(input_tokens), input_bpe_tokens)