codingutf8 2023 the huggingface team inc 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 orglicenseslicense2 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 refer to the full test suite in optimum library https github comhuggingfaceoptimumtreemaintestsbettertransformer modelid hfinternaltestingtinyrandomt5 tokenizer autotokenizer frompretrainedmodelid model automodelforseq2seqlm frompretrainedmodelid inp tokenizerthis is me returntensorspt model model tobettertransformer self asserttrueanybettertransformer in mod class name for mod in model namedmodules output model generateinp model model reversebettertransformer self assertfalseanybettertransformer in mod class name for mod in model namedmodules with tempfile temporarydirectory as tmpdirname model savepretrainedtmpdirname modelreloaded automodelforseq2seqlm frompretrainedtmpdirname self assertfalse anybettertransformer in mod class name for mod in modelreloaded namedmodules outputfrompretrained modelreloaded generateinp self asserttruetorch allcloseoutput outputfrompretrained def testerrorsavepretrainedself r the savepretrained method should raise a valueerror if the model is in bettertransformer mode all should be good if the model is reversed coding utf 8 2023 the huggingface team inc 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 refer to the full test suite in optimum library https github com huggingface optimum tree main tests bettertransformer classic tests to simply check if the conversion has been successfull the save_pretrained method should raise a valueerror if the model is in bettertransformer mode all should be good if the model is reversed
import tempfile
import unittest

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from transformers.testing_utils import (
    is_torch_available,
    require_optimum,
    require_torch,
    slow,
)


if is_torch_available():
    import torch


@require_torch
@require_optimum
@slow
class BetterTransformerIntegrationTest(unittest.TestCase):
    
    

    def test_transform_and_reverse(self):
        r
        model_id = "hf-internal-testing/tiny-random-t5"
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_id)

        inp = tokenizer("This is me", return_tensors="pt")

        model = model.to_bettertransformer()

        self.assertTrue(any("BetterTransformer" in mod.__class__.__name__ for _, mod in model.named_modules()))

        output = model.generate(**inp)

        model = model.reverse_bettertransformer()

        self.assertFalse(any("BetterTransformer" in mod.__class__.__name__ for _, mod in model.named_modules()))

        with tempfile.TemporaryDirectory() as tmpdirname:
            model.save_pretrained(tmpdirname)

            model_reloaded = AutoModelForSeq2SeqLM.from_pretrained(tmpdirname)

            self.assertFalse(
                any("BetterTransformer" in mod.__class__.__name__ for _, mod in model_reloaded.named_modules())
            )

            output_from_pretrained = model_reloaded.generate(**inp)
            self.assertTrue(torch.allclose(output, output_from_pretrained))

    def test_error_save_pretrained(self):
        r
        model_id = "hf-internal-testing/tiny-random-t5"
        model = AutoModelForSeq2SeqLM.from_pretrained(model_id)

        model = model.to_bettertransformer()

        with tempfile.TemporaryDirectory() as tmpdirname:
            with self.assertRaises(ValueError):
                model.save_pretrained(tmpdirname)

            model = model.reverse_bettertransformer()
            model.save_pretrained(tmpdirname)