Introduce code from Rust monorepo. This is a mess

2026-10-28 09:08:54 +00:00 · 2022-02-28 05:43:54 -05:00
parent 921fb49f9a
commit ba326bfcfd
1058 changed files with 192192 additions and 0 deletions
@@ -0,0 +1,86 @@
+# 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 os
+import tempfile
+
+import pytest
+import torch
+
+import nemo.collections.nlp as nemo_nlp
+from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
+from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer
+
+
+def do_export(model, name: str):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # Generate filename in the temporary directory.
+        tmp_file_name = os.path.join(tmpdir, name + '.onnx')
+        # Test export.
+        if torch.cuda.is_available():
+            model = model.cuda()
+        model.export(tmp_file_name)
+
+
+class TestHuggingFace:
+    @pytest.mark.unit
+    def test_list_pretrained_models(self):
+        pretrained_lm_models = nemo_nlp.modules.get_pretrained_lm_models_list()
+        assert len(pretrained_lm_models) > 0
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_get_pretrained_bert_model(self):
+        model = nemo_nlp.modules.get_lm_model(pretrained_model_name='bert-base-uncased')
+        assert isinstance(model, nemo_nlp.modules.BertEncoder)
+        do_export(model, "bert-base-uncased")
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_get_pretrained_distilbert_model(self):
+        model = nemo_nlp.modules.get_lm_model(pretrained_model_name='distilbert-base-uncased')
+        assert isinstance(model, nemo_nlp.modules.DistilBertEncoder)
+        do_export(model, "distilbert-base-uncased")
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_get_pretrained_roberta_model(self):
+        model = nemo_nlp.modules.get_lm_model(pretrained_model_name='roberta-base')
+        assert isinstance(model, nemo_nlp.modules.RobertaEncoder)
+        do_export(model, "roberta-base-uncased")
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_get_pretrained_albert_model(self):
+        model = nemo_nlp.modules.get_lm_model(pretrained_model_name='albert-base-v1')
+        assert isinstance(model, nemo_nlp.modules.AlbertEncoder)
+        do_export(model, "albert-base-v1")
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_get_pretrained_chinese_bert_wwm_model(self):
+        model_name = 'hfl/chinese-bert-wwm'
+        model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
+        assert isinstance(model, nemo_nlp.modules.BertModule)
+        tokenizer = get_tokenizer(tokenizer_name=model_name)
+        assert isinstance(tokenizer, AutoTokenizer)
+
+    @pytest.mark.with_downloads()
+    @pytest.mark.unit
+    def test_get_pretrained_arabic_model(self):
+        model_name = 'asafaya/bert-base-arabic'
+        model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name)
+        assert isinstance(model, nemo_nlp.modules.BertModule)
+        tokenizer = get_tokenizer(tokenizer_name=model_name)
+        assert isinstance(tokenizer, AutoTokenizer)