Hugging Face Language Models Implementation

art/estimators/__init__.py

@@ -18,6 +18,7 @@
 from art.estimators import classification
 from art.estimators import encoding
 from art.estimators import generation
+from art.estimators import language_modeling
 from art.estimators import object_detection
 from art.estimators import poison_mitigation
 from art.estimators import regression

art/estimators/language_modeling/__init__.py

@@ -0,0 +1,6 @@
+"""
+Module containing estimators for language modeling.
+"""
+from art.estimators.language_modeling.language_model import LanguageModelMixin
+
+from art.estimators.language_modeling.hugging_face import HuggingFaceLanguageModel

art/estimators/language_modeling/language_model.py

@@ -0,0 +1,135 @@
+# MIT License
+#
+# Copyright (C) The Adversarial Robustness Toolbox (ART) Authors 2023
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+# persons to whom the Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
+# Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+"""
+This module implements mixin abstract base class for all language models in ART.
+"""
+
+import abc
+from typing import Any
+
+from art.estimators.estimator import BaseEstimator
+
+
+class LanguageModelMixin(abc.ABC):
+    """
+    Mix-in base class for ART language models.
+    """
+
+    pass
+
+
+class LanguageModel(LanguageModelMixin, BaseEstimator, abc.ABC):
+    """
+    Abstract base class for ART language models used to define common types and methods.
+    """
+
+    @abc.abstractmethod
+    def tokenize(self, x: Any, **kwargs) -> Any:
+        """
+        Token the input `x`.
+
+        :param text: Samples to be tokenized.
+        :type text: Format as expected by the `tokenizer`
+        :return: Tokenized output by the tokenizer.
+        :rtype: Format as produced by the `tokenizer`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def encode(self, x: Any, **kwargs) -> Any:
+        """
+        Encode the input `x`.
+
+        :param text: Samples to be encoded.
+        :type text: Format as expected by the `tokenizer`
+        :return: Encoded output by the tokenizer.
+        :rtype: Format as produced by the `tokenizer`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def batch_encode(self, x: Any, **kwargs) -> Any:
+        """
+        Encode the input `x`.
+
+        :param text: Samples to be encoded.
+        :type text: Format as expected by the `tokenizer`
+        :return: Encoded output by the tokenizer.
+        :rtype: Format as produced by the `tokenizer`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def decode(self, x: Any, **kwargs) -> Any:
+        """
+        Decode the input `x`.
+
+        :param tokens: Samples to be decoded.
+        :type tokens: Format as expected by the `tokenizer`
+        :return: decoded output by the tokenizer.
+        :rtype: Format as produced by the `tokenizer`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def batch_decode(self, x: Any, **kwargs) -> Any:
+        """
+        Decode the input `x`.
+
+        :param tokens: Samples to be decoded.
+        :type tokens: Format as expected by the `tokenizer`
+        :return: decoded output by the tokenizer.
+        :rtype: Format as produced by the `tokenizer`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def predict(self, x: Any, **kwargs) -> Any:
+        """
+        Perform prediction of the language model for input `x`.
+
+        :param text: Samples to be tokenized.
+        :type text: Format as expected by the `tokenizer`
+        :return: Predictions by the model.
+        :rtype: Format as produced by the `model`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def generate(self, x: Any, **kwargs) -> Any:
+        """
+        Generate text using the language model from input `x`.
+
+        :param text: Samples to be tokenized.
+        :type text: Format as expected by the `tokenizer`
+        :return: Generated text by the model.
+        :rtype: Format as produced by the `tokenizer`
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def fit(self, x: Any, y: Any, **kwargs) -> None:
+        """
+        Fit the estimator using the training data `(x, y)`.
+
+        :param x: Training data.
+        :type x: Format as expected by the `model`
+        :param y: Target values.
+        :type y: Format as expected by the `model`
+        """
+        raise NotImplementedError

tests/estimators/language_modeling/conftest.py

@@ -0,0 +1,62 @@
+# MIT License
+#
+# Copyright (C) The Adversarial Robustness Toolbox (ART) Authors 2023
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+# persons to whom the Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
+# Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+
+import numpy as np
+import pytest
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.fixture()
+def get_text_data():
+    """
+    Sample text data to test language model estimators
+    """
+    text = [
+        "The capital of France is",
+        "The capital of Germany is",
+    ]
+    tokens = np.asarray([[464, 3139, 286, 4881, 318], [464, 3139, 286, 4486, 318]])
+    return text, tokens
+
+
+@pytest.fixture()
+def get_hugging_face_language_model(get_text_data):
+    """
+    This class tests the HuggingFaceLanguageModel estimator.
+    """
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    from art.estimators.language_modeling import HuggingFaceLanguageModel
+
+    # Define language model
+    model = AutoModelForCausalLM.from_pretrained("gpt2")
+    tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizer.pad_token = tokenizer.eos_token
+
+    language_model = HuggingFaceLanguageModel(
+        model=model,
+        tokenizer=tokenizer,
+    )
+
+    text, tokens = get_text_data
+
+    yield language_model, text, tokens

tests/estimators/language_modeling/test_hugging_face.py

@@ -0,0 +1,159 @@
+# MIT License
+#
+# Copyright (C) The Adversarial Robustness Toolbox (ART) Authors 2023
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
+# documentation files (the "Software"), to deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit
+# persons to whom the Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all copies or substantial portions of the
+# Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
+# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
+# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+from __future__ import absolute_import, division, print_function, unicode_literals
+
+import logging
+
+import numpy as np
+import pytest
+
+from tests.utils import ARTTestException
+
+logger = logging.getLogger(__name__)
+
+
+@pytest.mark.only_with_platform("huggingface")
+def test_tokenize(art_warning, get_hugging_face_language_model):
+    try:
+        language_model, text, _ = get_hugging_face_language_model
+
+        result = language_model.tokenize(x=text, padding=True, return_tensors="np")
+
+        assert list(result.keys()) == ["input_ids", "attention_mask"]
+
+        assert result["input_ids"].shape == (2, 5)
+        expected_input_ids = np.asarray([[464, 3139, 286, 4881, 318], [464, 3139, 286, 4486, 318]])
+        np.testing.assert_array_equal(result["input_ids"], expected_input_ids)
+
+        assert result["attention_mask"].shape == (2, 5)
+        expected_attention_masks = np.asarray([[1, 1, 1, 1, 1], [1, 1, 1, 1, 1]])
+        np.testing.assert_array_equal(result["attention_mask"], expected_attention_masks)
+
+    except ARTTestException as e:
+        art_warning(e)
+
+
+@pytest.mark.only_with_platform("huggingface")
+def test_encode(art_warning, get_hugging_face_language_model):
+    try:
+        language_model, text, tokens = get_hugging_face_language_model
+
+        result = language_model.encode(x=text[0])
+        np.testing.assert_array_equal(result, tokens[0])
+
+    except ARTTestException as e:
+        art_warning(e)
+
+
+@pytest.mark.only_with_platform("huggingface")
+def test_batch_encode(art_warning, get_hugging_face_language_model):
+    try:
+        language_model, text, tokens = get_hugging_face_language_model
+
+        result = language_model.batch_encode(x=text)
+        np.testing.assert_array_equal(result, tokens)
+
+    except ARTTestException as e:
+        art_warning(e)
+
+
+@pytest.mark.only_with_platform("huggingface")
+def test_decode(art_warning, get_hugging_face_language_model):
+    try:
+        language_model, text, tokens = get_hugging_face_language_model
+
+        result = language_model.decode(x=tokens[0])
+        np.testing.assert_array_equal(result, text[0])
+
+    except ARTTestException as e:
+        art_warning(e)
+
+
+@pytest.mark.only_with_platform("huggingface")
+def test_batch_decode(art_warning, get_hugging_face_language_model):
+    try:
+        language_model, text, tokens = get_hugging_face_language_model
+
+        result = language_model.batch_decode(x=tokens)
+        np.testing.assert_array_equal(result, text)
+
+    except ARTTestException as e:
+        art_warning(e)
+
+
+@pytest.mark.only_with_platform("huggingface")
+def test_predict(art_warning, get_hugging_face_language_model):
+    try:
+        language_model, text, _ = get_hugging_face_language_model
+
+        result = language_model.predict(x=text)
+
+        assert list(result.keys()) == ["logits", "past_key_values"]
+
+        assert result["logits"].shape == (2, 5, 50257)
+        expected_logits = np.asarray(
+            [
+                -36.287148,
+                -35.011105,
+                -38.079082,
+                -37.782597,
+                -36.693443,
+                -37.765358,
+                -35.628418,
+                -36.034225,
+                -34.752907,
+                -37.38435,
+            ]
+        )
+        np.testing.assert_array_almost_equal(result["logits"][0, 0, :10], expected_logits, decimal=2)
+
+        assert result["past_key_values"][0][0].shape == (2, 12, 5, 64)
+        expected_past_key_values = np.asarray(
+            [
+                -0.9420212,
+                1.9022521,
+                0.87219113,
+                -0.15018058,
+                0.4039965,
+                0.20760098,
+                -0.1169314,
+                0.6232018,
+                -1.7421292,
+                0.63497376,
+            ]
+        )
+        np.testing.assert_array_almost_equal(
+            result["past_key_values"][0][0][0, 0, 0, :10], expected_past_key_values, decimal=2
+        )
+
+    except ARTTestException as e:
+        art_warning(e)
+
+
+@pytest.mark.only_with_platform("huggingface")
+def test_generate(art_warning, get_hugging_face_language_model):
+    try:
+        language_model, text, _ = get_hugging_face_language_model
+
+        result = language_model.generate(x=text, max_new_tokens=1)
+
+        assert len(result[0]) > len(text[0])
+        assert len(result[1]) > len(text[1])
+
+    except ARTTestException as e:
+        art_warning(e)