# Copyright 2021 The TensorFlow Authors. All Rights Reserved.
#
# 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.

"""Test Subtokenizer and string helper methods."""

import collections
import tempfile

import tensorflow as tf

from official.nlp.transformer.utils import tokenizer


class SubtokenizerTest(tf.test.TestCase):

  def _init_subtokenizer(self, vocab_list):
    temp_file = tempfile.NamedTemporaryFile(delete=False)
    with tf.io.gfile.GFile(temp_file.name, "w") as w:
      for subtoken in vocab_list:
        w.write("'%s'" % subtoken)
        w.write("\n")
    return tokenizer.Subtokenizer(temp_file.name, reserved_tokens=[])

  def test_encode(self):
    vocab_list = ["123_", "test", "ing_"]
    subtokenizer = self._init_subtokenizer(vocab_list)
    s = "testing 123"
    encoded_list = subtokenizer.encode(s)
    self.assertEqual([1, 2, 0], encoded_list)

  def test_decode(self):
    vocab_list = ["123_", "test", "ing_"]
    subtokenizer = self._init_subtokenizer(vocab_list)
    encoded_list = [1, 2, 0]  # testing 123
    decoded_str = subtokenizer.decode(encoded_list)
    self.assertEqual("testing 123", decoded_str)

  def test_subtoken_ids_to_tokens(self):
    vocab_list = ["123_", "test", "ing_"]
    subtokenizer = self._init_subtokenizer(vocab_list)
    encoded_list = [1, 2, 0]  # testing 123
    token_list = subtokenizer._subtoken_ids_to_tokens(encoded_list)
    self.assertEqual([u"testing", u"123"], token_list)


class StringHelperTest(tf.test.TestCase):

  def test_split_string_to_tokens(self):
    text = "test? testing 123."

    tokens = tokenizer._split_string_to_tokens(text,
                                               tokenizer._ALPHANUMERIC_CHAR_SET)
    self.assertEqual(["test", "? ", "testing", "123", "."], tokens)

  def test_join_tokens_to_string(self):
    tokens = ["test", "? ", "testing", "123", "."]

    s = tokenizer._join_tokens_to_string(tokens,
                                         tokenizer._ALPHANUMERIC_CHAR_SET)
    self.assertEqual("test? testing 123.", s)

  def test_escape_token(self):
    token = u"abc_\\4"
    alphabet = set("abc_\\u;")

    escaped_token = tokenizer._escape_token(token, alphabet)
    self.assertEqual("abc\\u\\\\\\52;_", escaped_token)

  def test_unescape_token(self):
    escaped_token = u"Underline: \\u, Backslash: \\\\, Unicode: \\52;"

    unescaped_token = tokenizer._unescape_token(escaped_token)
    self.assertEqual("Underline: _, Backslash: \\, Unicode: 4", unescaped_token)

  def test_list_to_index_dict(self):
    lst = ["test", "strings"]

    d = tokenizer._list_to_index_dict(lst)
    self.assertDictEqual({"test": 0, "strings": 1}, d)

  def test_split_token_to_subtokens(self):
    token = "abc"
    subtoken_dict = {"a": 0, "b": 1, "c": 2, "ab": 3}
    max_subtoken_length = 2

    subtokens = tokenizer._split_token_to_subtokens(token, subtoken_dict,
                                                    max_subtoken_length)
    self.assertEqual(["ab", "c"], subtokens)

  def test_generate_alphabet_dict(self):
    s = ["testing", "123"]
    reserved_tokens = ["???"]

    alphabet = tokenizer._generate_alphabet_dict(s, reserved_tokens)
    self.assertIn("?", alphabet)
    self.assertIn("t", alphabet)
    self.assertIn("e", alphabet)
    self.assertIn("s", alphabet)
    self.assertIn("i", alphabet)
    self.assertIn("n", alphabet)
    self.assertIn("g", alphabet)
    self.assertIn("1", alphabet)
    self.assertIn("2", alphabet)
    self.assertIn("3", alphabet)

  def test_count_and_gen_subtokens(self):
    token_counts = {"abc": 5}
    alphabet = set("abc_")
    subtoken_dict = {"a": 0, "b": 1, "c": 2, "_": 3}
    max_subtoken_length = 2

    subtoken_counts = tokenizer._count_and_gen_subtokens(
        token_counts, alphabet, subtoken_dict, max_subtoken_length)

    self.assertIsInstance(subtoken_counts, collections.defaultdict)
    self.assertDictEqual(
        {
            "a": 5,
            "b": 5,
            "c": 5,
            "_": 5,
            "ab": 5,
            "bc": 5,
            "c_": 5,
            "abc": 5,
            "bc_": 5,
            "abc_": 5
        }, subtoken_counts)

  def test_filter_and_bucket_subtokens(self):
    subtoken_counts = collections.defaultdict(int, {
        "a": 2,
        "b": 4,
        "c": 1,
        "ab": 6,
        "ac": 3,
        "abbc": 5
    })
    min_count = 3

    subtoken_buckets = tokenizer._filter_and_bucket_subtokens(
        subtoken_counts, min_count)

    self.assertEqual(len(subtoken_buckets[0]), 0)
    self.assertEqual(set("b"), subtoken_buckets[1])
    self.assertEqual(set(["ab", "ac"]), subtoken_buckets[2])
    self.assertEqual(len(subtoken_buckets[3]), 0)
    self.assertEqual(set(["abbc"]), subtoken_buckets[4])

  def test_gen_new_subtoken_list(self):
    subtoken_counts = collections.defaultdict(int, {
        "translate": 10,
        "t": 40,
        "tr": 16,
        "tra": 12
    })
    min_count = 5
    alphabet = set("translate")
    reserved_tokens = ["reserved", "tokens"]

    subtoken_list, max_token_length = tokenizer._gen_new_subtoken_list(
        subtoken_counts, min_count, alphabet, reserved_tokens)

    # Check that "tra" isn"t in the list (its count should be decremented to 2,
    # so it should not be added to the canddiate list).
    self.assertNotIn("tra", subtoken_list)

    self.assertIn("tr", subtoken_list)
    self.assertIn("t", subtoken_list)

    self.assertEqual(len("translate"), max_token_length)

  def test_generate_subtokens(self):
    token_counts = {"ab": 1, "bc": 3, "abc": 5}
    alphabet = set("abc_")
    min_count = 100
    num_iterations = 1
    reserved_tokens = ["reserved", "tokens"]

    vocab_list = tokenizer._generate_subtokens(token_counts, alphabet,
                                               min_count, num_iterations,
                                               reserved_tokens)

    # Check that reserved tokens are at the front of the list
    self.assertEqual(vocab_list[:2], reserved_tokens)

    # Check that each character in alphabet is in the vocab list
    for c in alphabet:
      self.assertIn(c, vocab_list)


if __name__ == "__main__":
  tf.test.main()

[((6843, 6857), 'tensorflow.test.main', 'tf.test.main', ([], {}), '()\n', (6855, 6857), True, 'import tensorflow as tf\n'), ((907, 948), 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', ([], {'delete': '(False)'}), '(delete=False)\n', (934, 948), False, 'import tempfile\n'), ((1110, 1168), 'official.nlp.transformer.utils.tokenizer.Subtokenizer', 'tokenizer.Subtokenizer', (['temp_file.name'], {'reserved_tokens': '[]'}), '(temp_file.name, reserved_tokens=[])\n', (1132, 1168), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((2137, 2210), 'official.nlp.transformer.utils.tokenizer._split_string_to_tokens', 'tokenizer._split_string_to_tokens', (['text', 'tokenizer._ALPHANUMERIC_CHAR_SET'], {}), '(text, tokenizer._ALPHANUMERIC_CHAR_SET)\n', (2170, 2210), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((2434, 2508), 'official.nlp.transformer.utils.tokenizer._join_tokens_to_string', 'tokenizer._join_tokens_to_string', (['tokens', 'tokenizer._ALPHANUMERIC_CHAR_SET'], {}), '(tokens, tokenizer._ALPHANUMERIC_CHAR_SET)\n', (2466, 2508), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((2711, 2751), 'official.nlp.transformer.utils.tokenizer._escape_token', 'tokenizer._escape_token', (['token', 'alphabet'], {}), '(token, alphabet)\n', (2734, 2751), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((2942, 2982), 'official.nlp.transformer.utils.tokenizer._unescape_token', 'tokenizer._unescape_token', (['escaped_token'], {}), '(escaped_token)\n', (2967, 2982), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((3147, 3181), 'official.nlp.transformer.utils.tokenizer._list_to_index_dict', 'tokenizer._list_to_index_dict', (['lst'], {}), '(lst)\n', (3176, 3181), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((3406, 3484), 'official.nlp.transformer.utils.tokenizer._split_token_to_subtokens', 'tokenizer._split_token_to_subtokens', (['token', 'subtoken_dict', 'max_subtoken_length'], {}), '(token, subtoken_dict, max_subtoken_length)\n', (3441, 3484), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((3705, 3758), 'official.nlp.transformer.utils.tokenizer._generate_alphabet_dict', 'tokenizer._generate_alphabet_dict', (['s', 'reserved_tokens'], {}), '(s, reserved_tokens)\n', (3738, 3758), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((4311, 4409), 'official.nlp.transformer.utils.tokenizer._count_and_gen_subtokens', 'tokenizer._count_and_gen_subtokens', (['token_counts', 'alphabet', 'subtoken_dict', 'max_subtoken_length'], {}), '(token_counts, alphabet, subtoken_dict,\n    max_subtoken_length)\n', (4345, 4409), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((4845, 4932), 'collections.defaultdict', 'collections.defaultdict', (['int', "{'a': 2, 'b': 4, 'c': 1, 'ab': 6, 'ac': 3, 'abbc': 5}"], {}), "(int, {'a': 2, 'b': 4, 'c': 1, 'ab': 6, 'ac': 3,\n    'abbc': 5})\n", (4868, 4932), False, 'import collections\n'), ((5035, 5101), 'official.nlp.transformer.utils.tokenizer._filter_and_bucket_subtokens', 'tokenizer._filter_and_bucket_subtokens', (['subtoken_counts', 'min_count'], {}), '(subtoken_counts, min_count)\n', (5073, 5101), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((5455, 5532), 'collections.defaultdict', 'collections.defaultdict', (['int', "{'translate': 10, 't': 40, 'tr': 16, 'tra': 12}"], {}), "(int, {'translate': 10, 't': 40, 'tr': 16, 'tra': 12})\n", (5478, 5532), False, 'import collections\n'), ((5715, 5806), 'official.nlp.transformer.utils.tokenizer._gen_new_subtoken_list', 'tokenizer._gen_new_subtoken_list', (['subtoken_counts', 'min_count', 'alphabet', 'reserved_tokens'], {}), '(subtoken_counts, min_count, alphabet,\n    reserved_tokens)\n', (5747, 5806), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((6366, 6467), 'official.nlp.transformer.utils.tokenizer._generate_subtokens', 'tokenizer._generate_subtokens', (['token_counts', 'alphabet', 'min_count', 'num_iterations', 'reserved_tokens'], {}), '(token_counts, alphabet, min_count,\n    num_iterations, reserved_tokens)\n', (6395, 6467), False, 'from official.nlp.transformer.utils import tokenizer\n'), ((959, 997), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', (['temp_file.name', '"""w"""'], {}), "(temp_file.name, 'w')\n", (976, 997), True, 'import tensorflow as tf\n')]