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

23
24
25
26
27
from transformers import BatchEncoding, MarianTokenizer
from transformers.testing_utils import _sentencepiece_available, _torch_available, require_sentencepiece


if _sentencepiece_available:
Sylvain Gugger's avatar
Sylvain Gugger committed
28
    from transformers.models.marian.tokenization_marian import save_json, vocab_files_names
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
FRAMEWORK = "pt" if _torch_available else "tf"
39
40


41
@require_sentencepiece
42
43
44
class MarianTokenizationTest(TokenizerTesterMixin, unittest.TestCase):

    tokenizer_class = MarianTokenizer
45
    test_rust_tokenizer = False
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

    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)
        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"])

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

61
62
    def get_tokenizer(self, **kwargs) -> MarianTokenizer:
        return MarianTokenizer.from_pretrained(self.tmpdirname, **kwargs)
63

64
    def get_input_output_texts(self, tokenizer):
65
66
67
68
69
70
71
        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")
72
        batch = en_de_tokenizer.prepare_seq2seq_batch(["I am a small frog"], return_tensors=None)
73
74
75
        self.assertIsInstance(batch, BatchEncoding)
        expected = [38, 121, 14, 697, 38848, 0]
        self.assertListEqual(expected, batch.input_ids[0])
76
77
78
79
80
81

        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)
82
83
84
85

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

86
        batch = tok.prepare_seq2seq_batch(["I am a small frog" * 1000, "I am a small frog"], return_tensors=FRAMEWORK)
87
88
89
90
91
        self.assertIsInstance(batch, BatchEncoding)
        self.assertEqual(batch.input_ids.shape, (2, 512))

    def test_outputs_can_be_shorter(self):
        tok = self.get_tokenizer()
92
        batch_smaller = tok.prepare_seq2seq_batch(["I am a tiny frog", "I am a small frog"], return_tensors=FRAMEWORK)
93
94
        self.assertIsInstance(batch_smaller, BatchEncoding)
        self.assertEqual(batch_smaller.input_ids.shape, (2, 10))