






"""simple docstring"""



import os
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers.models.realm.configuration_realm import RealmConfig
from transformers.models.realm.retrieval_realm import _REALM_BLOCK_RECORDS_FILENAME, RealmRetriever
from transformers.models.realm.tokenization_realm import VOCAB_FILES_NAMES, RealmTokenizer




class 			lowerCAmelCase__			(    __magic_name__   ):
   def       __a							(		self      :	Any							):


        '''simple docstring'''
        UpperCAmelCase__						:    Tuple							=							tempfile.mkdtemp()
        UpperCAmelCase__						:    Optional[Any]							=							5

        # Realm tok
        UpperCAmelCase__						:    Any							=							[
            "[UNK]",
            "[CLS]",
            "[SEP]",
            "[PAD]",
            "[MASK]",
            "test",
            "question",
            "this",
            "is",
            "the",
            "first",
            "second",
            "third",
            "fourth",
            "fifth",
            "record",
            "want",
            "##want",
            "##ed",
            "wa",
            "un",
            "runn",
            "##ing",
            ",",
            "low",
            "lowest",
        ]
        UpperCAmelCase__						:    str							=							os.path.join(self.tmpdirname							,			"realm_tokenizer"							)
        os.makedirs(snake_case__							,			exist_ok=snake_case__							)
        UpperCAmelCase__						:    str							=							os.path.join(snake_case__							,			VOCAB_FILES_NAMES["vocab_file"]							)
        with open(self.vocab_file							,			"w"							,			encoding="utf-8"							) as vocab_writer:
             vocab_writer.write("".join([x + "\n" for x in vocab_tokens]							)							)

        UpperCAmelCase__						:    str							=							os.path.join(self.tmpdirname							,			"realm_block_records"							)
        os.makedirs(snake_case__							,			exist_ok=snake_case__							)
   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname							,			"realm_tokenizer"							)							)
   def       __a							(		self      :	str							):


        '''simple docstring'''
        shutil.rmtree(self.tmpdirname							)
   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[Any]							=							RealmConfig(num_block_records=self.num_block_records							)
        return config
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[int]							=							Dataset.from_dict(
            {
                "id": ["0", "1"],
                "question": ["foo", "bar"],
                "answers": [["Foo", "Bar"], ["Bar"]],
            }							)
        return dataset
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Tuple							=							np.array(
            [
                B"This is the first record",
                B"This is the second record",
                B"This is the third record",
                B"This is the fourth record",
                B"This is the fifth record",
                B"This is a longer longer longer record",
            ]							,			dtype=snake_case__							,			)
        return block_records
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Tuple							=							RealmRetriever(
            block_records=self.get_dummy_block_records()							,			tokenizer=self.get_tokenizer()							,			)
        return retriever
   def       __a							(		self      :	Optional[int]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Dict							=							self.get_config()
        UpperCAmelCase__						:    List[str]							=							self.get_dummy_retriever()
        UpperCAmelCase__						:    int							=							retriever.tokenizer

        UpperCAmelCase__						:    Tuple							=							np.array([0, 3]							,			dtype="long"							)
        UpperCAmelCase__						:    Union[str, Any]							=							tokenizer(["Test question"]							).input_ids
        UpperCAmelCase__						:    Union[str, Any]							=							tokenizer(
            ["the fourth"]							,			add_special_tokens=snake_case__							,			return_token_type_ids=snake_case__							,			return_attention_mask=snake_case__							,			).input_ids
        UpperCAmelCase__						:    Optional[Any]							=							config.reader_seq_len

        UpperCAmelCase__							,				UpperCAmelCase__							,				UpperCAmelCase__							,				UpperCAmelCase__						:    List[Any]							=							retriever(
            snake_case__							,			snake_case__							,			answer_ids=snake_case__							,			max_length=snake_case__							,			return_tensors="np"							)

        self.assertEqual(len(snake_case__							)							,			2							)
        self.assertEqual(len(snake_case__							)							,			2							)
        self.assertEqual(len(snake_case__							)							,			2							)
        self.assertEqual(concat_inputs.input_ids.shape							,			(2, 1_0)							)
        self.assertEqual(concat_inputs.attention_mask.shape							,			(2, 1_0)							)
        self.assertEqual(concat_inputs.token_type_ids.shape							,			(2, 1_0)							)
        self.assertEqual(concat_inputs.special_tokens_mask.shape							,			(2, 1_0)							)
        self.assertEqual(
            tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[0]							)							,			["[CLS]", "test", "question", "[SEP]", "this", "is", "the", "first", "record", "[SEP]"]							,			)
        self.assertEqual(
            tokenizer.convert_ids_to_tokens(concat_inputs.input_ids[1]							)							,			["[CLS]", "test", "question", "[SEP]", "this", "is", "the", "fourth", "record", "[SEP]"]							,			)
   def       __a							(		self      :	List[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[int]							=							self.get_config()
        UpperCAmelCase__						:    Any							=							self.get_dummy_retriever()
        UpperCAmelCase__						:    List[str]							=							retriever.tokenizer

        UpperCAmelCase__						:    List[Any]							=							np.array([0, 3, 5]							,			dtype="long"							)
        UpperCAmelCase__						:    int							=							tokenizer(["Test question"]							).input_ids
        UpperCAmelCase__						:    Tuple							=							tokenizer(
            ["the fourth", "longer longer"]							,			add_special_tokens=snake_case__							,			return_token_type_ids=snake_case__							,			return_attention_mask=snake_case__							,			).input_ids
        UpperCAmelCase__						:    Any							=							config.reader_seq_len

        UpperCAmelCase__							,				UpperCAmelCase__							,				UpperCAmelCase__							,				UpperCAmelCase__						:    Dict							=							retriever(
            snake_case__							,			snake_case__							,			answer_ids=snake_case__							,			max_length=snake_case__							,			return_tensors="np"							)

        self.assertEqual([False, True, True]							,			snake_case__							)
        self.assertEqual([[-1, -1, -1], [6, -1, -1], [6, 7, 8]]							,			snake_case__							)
        self.assertEqual([[-1, -1, -1], [7, -1, -1], [7, 8, 9]]							,			snake_case__							)




   def       __a							(		self      :	Optional[Any]							):


        '''simple docstring'''
        UpperCAmelCase__						:    str							=							self.get_dummy_retriever()
        retriever.save_pretrained(os.path.join(self.tmpdirname							,			"realm_block_records"							)							)

        # Test local path
        UpperCAmelCase__						:    Tuple							=							retriever.from_pretrained(os.path.join(self.tmpdirname							,			"realm_block_records"							)							)
        self.assertEqual(retriever.block_records[0]							,			B"This is the first record"							)

        # Test mocked remote path
        with patch("transformers.models.realm.retrieval_realm.hf_hub_download"							) as mock_hf_hub_download:
             UpperCAmelCase__						:    List[str]							=							os.path.join(
                 os.path.join(self.tmpdirname							,			"realm_block_records"							)							,			_REALM_BLOCK_RECORDS_FILENAME							)
             UpperCAmelCase__						:    Tuple							=							RealmRetriever.from_pretrained("google/realm-cc-news-pretrained-openqa"							)

        self.assertEqual(retriever.block_records[0]							,			B"This is the first record"							)










"""simple docstring"""



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, FlaxMTaForConditionalGeneration
 from transformers.models.ta.modeling_flax_ta import shift_tokens_right




@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class 			lowerCAmelCase__			(    unittest.TestCase   ):
   @slow
   def       __a							(		self      :	List[str]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Union[str, Any]							=							FlaxMTaForConditionalGeneration.from_pretrained("google/mt5-small"							)
        UpperCAmelCase__						:    str							=							AutoTokenizer.from_pretrained("google/mt5-small"							)

        UpperCAmelCase__						:    str							=							tokenizer("Hello there"							,			return_tensors="np"							).input_ids
        UpperCAmelCase__						:    Optional[Any]							=							tokenizer("Hi I am"							,			return_tensors="np"							).input_ids

        UpperCAmelCase__						:    Any							=							shift_tokens_right(snake_case__							,			model.config.pad_token_id							,			model.config.decoder_start_token_id							)

        UpperCAmelCase__						:    Any							=							model(snake_case__							,			decoder_input_ids=snake_case__							).logits
        UpperCAmelCase__						:    Any							=							optax.softmax_cross_entropy(snake_case__							,			onehot(snake_case__							,			logits.shape[-1]							)							).mean()

        UpperCAmelCase__						:    List[str]							=							-(labels.shape[-1] * loss.item())

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


