test_tokenization_marian.py 3.84 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# coding=utf-8
# Copyright 2020 Huggingface
#
# 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 os
17
import tempfile
18
19
20
21
import unittest
from pathlib import Path
from shutil import copyfile

22
from transformers import BatchEncoding, MarianTokenizer
23
from transformers.file_utils import is_sentencepiece_available, is_tf_available, is_torch_available
24
from transformers.testing_utils import require_sentencepiece
25
26


27
if is_sentencepiece_available():
28
    from transformers.models.marian.tokenization_marian import VOCAB_FILES_NAMES, save_json
29
30
31
32
33
34
35
36
37

from .test_tokenization_common import TokenizerTesterMixin


SAMPLE_SP = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fixtures/test_sentencepiece.model")

mock_tokenizer_config = {"target_lang": "fi", "source_lang": "en"}
zh_code = ">>zh<<"
ORG_NAME = "Helsinki-NLP/"
38
39
40
41
42
43
44

if is_torch_available():
    FRAMEWORK = "pt"
elif is_tf_available():
    FRAMEWORK = "tf"
else:
    FRAMEWORK = "jax"
45
46


47
@require_sentencepiece
48
49
50
class MarianTokenizationTest(TokenizerTesterMixin, unittest.TestCase):

    tokenizer_class = MarianTokenizer
51
    test_rust_tokenizer = False
52
    test_sentencepiece = True
53
54
55
56
57
58

    def setUp(self):
        super().setUp()
        vocab = ["</s>", "<unk>", "鈻乀his", "鈻乮s", "鈻乤", "鈻乼", "est", "\u0120", "<pad>"]
        vocab_tokens = dict(zip(vocab, range(len(vocab))))
        save_dir = Path(self.tmpdirname)
59
60
61
62
63
        save_json(vocab_tokens, save_dir / VOCAB_FILES_NAMES["vocab"])
        save_json(mock_tokenizer_config, save_dir / VOCAB_FILES_NAMES["tokenizer_config_file"])
        if not (save_dir / VOCAB_FILES_NAMES["source_spm"]).exists():
            copyfile(SAMPLE_SP, save_dir / VOCAB_FILES_NAMES["source_spm"])
            copyfile(SAMPLE_SP, save_dir / VOCAB_FILES_NAMES["target_spm"])
64
65
66
67

        tokenizer = MarianTokenizer.from_pretrained(self.tmpdirname)
        tokenizer.save_pretrained(self.tmpdirname)

68
69
    def get_tokenizer(self, **kwargs) -> MarianTokenizer:
        return MarianTokenizer.from_pretrained(self.tmpdirname, **kwargs)
70

71
    def get_input_output_texts(self, tokenizer):
72
73
74
75
76
77
78
        return (
            "This is a test",
            "This is a test",
        )

    def test_tokenizer_equivalence_en_de(self):
        en_de_tokenizer = MarianTokenizer.from_pretrained(f"{ORG_NAME}opus-mt-en-de")
79
        batch = en_de_tokenizer(["I am a small frog"], return_tensors=None)
80
81
82
        self.assertIsInstance(batch, BatchEncoding)
        expected = [38, 121, 14, 697, 38848, 0]
        self.assertListEqual(expected, batch.input_ids[0])
83
84
85
86
87
88

        save_dir = tempfile.mkdtemp()
        en_de_tokenizer.save_pretrained(save_dir)
        contents = [x.name for x in Path(save_dir).glob("*")]
        self.assertIn("source.spm", contents)
        MarianTokenizer.from_pretrained(save_dir)
89
90
91
92

    def test_outputs_not_longer_than_maxlen(self):
        tok = self.get_tokenizer()

93
94
95
        batch = tok(
            ["I am a small frog" * 1000, "I am a small frog"], padding=True, truncation=True, return_tensors=FRAMEWORK
        )
96
97
98
99
100
        self.assertIsInstance(batch, BatchEncoding)
        self.assertEqual(batch.input_ids.shape, (2, 512))

    def test_outputs_can_be_shorter(self):
        tok = self.get_tokenizer()
101
        batch_smaller = tok(["I am a tiny frog", "I am a small frog"], padding=True, return_tensors=FRAMEWORK)
102
103
        self.assertIsInstance(batch_smaller, BatchEncoding)
        self.assertEqual(batch_smaller.input_ids.shape, (2, 10))