2021 the huggingface team 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 for comparision run import t5 pip install t50 7 1 from t5 data sentencepiecevocabulary import sentencepiecevocabulary pathtomtfsmallmt5checkpoint fillin pathtomtfsmallmt5spmmodelpath fillin t5model t5 models mtfmodelmodeldirpathtomtfsmallmt5checkpoint batchsize1 tpunone vocab sentencepiecevocabularypathtomtfsmallmt5spmmodelpath score t5model scoreinputshello there targetshi i am vocabularyvocab 2021 the huggingface team 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 for comparision run import t5 pip install t5 0 7 1 from t5 data sentencepiece_vocabulary import sentencepiecevocabulary path_to_mtf_small_mt5_checkpoint fill_in path_to_mtf_small_mt5_spm_model_path fill_in t5_model t5 models mtfmodel model_dir path_to_mtf_small_mt5_checkpoint batch_size 1 tpu none vocab sentencepiecevocabulary path_to_mtf_small_mt5_spm_model_path score t5_model score inputs hello there targets hi i am vocabulary vocab
import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
    import optax
    from flax.training.common_utils import onehot

    from transformers import AutoTokenizer, FlaxMT5ForConditionalGeneration
    from transformers.models.t5.modeling_flax_t5 import shift_tokens_right


@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class MT5IntegrationTest(unittest.TestCase):
    @slow
    def test_small_integration_test(self):
        

        model = FlaxMT5ForConditionalGeneration.from_pretrained("google/mt5-small")
        tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")

        input_ids = tokenizer("Hello there", return_tensors="np").input_ids
        labels = tokenizer("Hi I am", return_tensors="np").input_ids

        decoder_input_ids = shift_tokens_right(labels, model.config.pad_token_id, model.config.decoder_start_token_id)

        logits = model(input_ids, decoder_input_ids=decoder_input_ids).logits
        loss = optax.softmax_cross_entropy(logits, onehot(labels, logits.shape[-1])).mean()

        mtf_score = -(labels.shape[-1] * loss.item())

        EXPECTED_SCORE = -84.9127
        self.assertTrue(abs(mtf_score - EXPECTED_SCORE) < 1e-4)