File size: 2,282 Bytes
7934b29 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | # Copyright (c) 2020, NVIDIA CORPORATION. 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.
import tempfile
import pytest
from nemo.collections.common.tokenizers.regex_tokenizer import RegExTokenizer
DEFAULT_REGEX = (
r"""\[[^\]]+]|Br?|Cl?|N|O|S|P|F|I|b|c|n|o|s|p|\(|\)|\.|=|#|-|\+|\\\\|\/|:|~|@|\?|>|\*|\$|\%[0-9]{2}|[0-9]"""
)
class TestRegexTokenizer:
def create_test_data(self):
data_file = tempfile.NamedTemporaryFile(mode='w+', delete=False)
data_file.writelines(
"""zinc_id,smiles
ZINC000510438538,FC(F)Oc1ccc([C@H](NCc2cnc3ccccn23)C(F)(F)F)cc1
"""
)
data_file_path = str(data_file.name)
data_file.close()
return data_file_path
@pytest.mark.unit
def test_create_vocab(self):
data_file_path = self.create_test_data()
tokenizer = RegExTokenizer(regex=DEFAULT_REGEX)
tokenizer.build_vocab_from_csv(data_csv_file=data_file_path)
assert len(tokenizer.vocab) == 18
@pytest.mark.unit
def test_text_2_tokens(self):
tokenizer = RegExTokenizer(regex=DEFAULT_REGEX)
tokens = tokenizer.text_to_tokens("Zc")
assert ''.join(tokens) == 'Zc'
@pytest.mark.unit
def test_text_2_ids(self):
data_file_path = self.create_test_data()
tokenizer = RegExTokenizer(regex=DEFAULT_REGEX)
tokenizer.build_vocab_from_csv(data_csv_file=data_file_path)
ids = tokenizer.text_to_ids("Zc")
assert ','.join(list(map(lambda x: str(x), ids))) == '1,11'
@pytest.mark.unit
def test_tokens_2_text(self):
tokenizer = RegExTokenizer(regex=DEFAULT_REGEX)
tokens = tokenizer.tokens_to_text(['^', 'Z', 'c', '&'])
assert ''.join(tokens) == 'Zc'
|