# 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()

[((204, 2, 204, 16), 'tensorflow.test.main', 'tf.test.main', ({}, {}), '()', True, 'import tensorflow as tf\n'), ((28, 16, 28, 57), 'tempfile.NamedTemporaryFile', 'tempfile.NamedTemporaryFile', (), '', False, 'import tempfile\n'), ((33, 11, 33, 69), 'official.nlp.transformer.utils.tokenizer.Subtokenizer', 'tokenizer.Subtokenizer', (), '', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((62, 13, 63, 80), 'official.nlp.transformer.utils.tokenizer._split_string_to_tokens', 'tokenizer._split_string_to_tokens', ({(62, 47, 62, 51): 'text', (63, 47, 63, 79): 'tokenizer._ALPHANUMERIC_CHAR_SET'}, {}), '(text, tokenizer._ALPHANUMERIC_CHAR_SET)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((69, 8, 70, 74), 'official.nlp.transformer.utils.tokenizer._join_tokens_to_string', 'tokenizer._join_tokens_to_string', ({(69, 41, 69, 47): 'tokens', (70, 41, 70, 73): 'tokenizer._ALPHANUMERIC_CHAR_SET'}, {}), '(tokens, tokenizer._ALPHANUMERIC_CHAR_SET)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((77, 20, 77, 60), 'official.nlp.transformer.utils.tokenizer._escape_token', 'tokenizer._escape_token', ({(77, 44, 77, 49): 'token', (77, 51, 77, 59): 'alphabet'}, {}), '(token, alphabet)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((83, 22, 83, 62), 'official.nlp.transformer.utils.tokenizer._unescape_token', 'tokenizer._unescape_token', ({(83, 48, 83, 61): 'escaped_token'}, {}), '(escaped_token)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((89, 8, 89, 42), 'official.nlp.transformer.utils.tokenizer._list_to_index_dict', 'tokenizer._list_to_index_dict', ({(89, 38, 89, 41): 'lst'}, {}), '(lst)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((97, 16, 98, 72), 'official.nlp.transformer.utils.tokenizer._split_token_to_subtokens', 'tokenizer._split_token_to_subtokens', ({(97, 52, 97, 57): 'token', (97, 59, 97, 72): 'subtoken_dict', (98, 52, 98, 71): 'max_subtoken_length'}, {}), '(token, subtoken_dict, max_subtoken_length)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((105, 15, 105, 68), 'official.nlp.transformer.utils.tokenizer._generate_alphabet_dict', 'tokenizer._generate_alphabet_dict', ({(105, 49, 105, 50): 's', (105, 52, 105, 67): 'reserved_tokens'}, {}), '(s, reserved_tokens)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((123, 22, 124, 67), 'official.nlp.transformer.utils.tokenizer._count_and_gen_subtokens', 'tokenizer._count_and_gen_subtokens', ({(124, 8, 124, 20): 'token_counts', (124, 22, 124, 30): 'alphabet', (124, 32, 124, 45): 'subtoken_dict', (124, 47, 124, 66): 'max_subtoken_length'}, {}), '(token_counts, alphabet, subtoken_dict,\n    max_subtoken_length)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((142, 22, 149, 6), 'collections.defaultdict', 'collections.defaultdict', ({(142, 46, 142, 49): 'int', (142, 51, 149, 5): "{'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})", False, 'import collections\n'), ((152, 23, 153, 35), 'official.nlp.transformer.utils.tokenizer._filter_and_bucket_subtokens', 'tokenizer._filter_and_bucket_subtokens', ({(153, 8, 153, 23): 'subtoken_counts', (153, 25, 153, 34): 'min_count'}, {}), '(subtoken_counts, min_count)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((162, 22, 167, 6), 'collections.defaultdict', 'collections.defaultdict', ({(162, 46, 162, 49): 'int', (162, 51, 167, 5): "{'translate': 10, 't': 40, 'tr': 16, 'tra': 12}"}, {}), "(int, {'translate': 10, 't': 40, 'tr': 16, 'tra': 12})", False, 'import collections\n'), ((172, 38, 173, 62), 'official.nlp.transformer.utils.tokenizer._gen_new_subtoken_list', 'tokenizer._gen_new_subtoken_list', ({(173, 8, 173, 23): 'subtoken_counts', (173, 25, 173, 34): 'min_count', (173, 36, 173, 44): 'alphabet', (173, 46, 173, 61): 'reserved_tokens'}, {}), '(subtoken_counts, min_count, alphabet,\n    reserved_tokens)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((191, 17, 193, 63), 'official.nlp.transformer.utils.tokenizer._generate_subtokens', 'tokenizer._generate_subtokens', ({(191, 47, 191, 59): 'token_counts', (191, 61, 191, 69): 'alphabet', (192, 47, 192, 56): 'min_count', (192, 58, 192, 72): 'num_iterations', (193, 47, 193, 62): 'reserved_tokens'}, {}), '(token_counts, alphabet, min_count,\n    num_iterations, reserved_tokens)', False, 'from official.nlp.transformer.utils import tokenizer\n'), ((29, 9, 29, 47), 'tensorflow.io.gfile.GFile', 'tf.io.gfile.GFile', ({(29, 27, 29, 41): 'temp_file.name', (29, 43, 29, 46): '"""w"""'}, {}), "(temp_file.name, 'w')", True, 'import tensorflow as tf\n')]