"""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  snake_case__      (       snake_case_       ):


 def 						a__   ( self						):
        __a       =    tempfile.mkdtemp()
        __a       =    5

        # Realm tok
        __a       =    [
            "[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",
        ]
        __a       =    os.path.join(self.tmpdirname    ,		"realm_tokenizer"						)
        os.makedirs(lowerCamelCase    ,		exist_ok=lowerCamelCase						)
        __a       =    os.path.join(lowerCamelCase    ,		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]						)						)

        __a       =    os.path.join(self.tmpdirname    ,		"realm_block_records"						)
        os.makedirs(lowerCamelCase    ,		exist_ok=lowerCamelCase						)


 def 						a__   ( self						):
        return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname    ,		"realm_tokenizer"						)						)


 def 						a__   ( self						):
        shutil.rmtree(self.tmpdirname						)


 def 						a__   ( self						):
        __a       =    RealmConfig(num_block_records=self.num_block_records						)
        return config


 def 						a__   ( self						):
        __a       =    Dataset.from_dict(
            {
                "id": ["0", "1"],
                "question": ["foo", "bar"],
                "answers": [["Foo", "Bar"], ["Bar"]],
            }						)
        return dataset


 def 						a__   ( self						):
        __a       =    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=lowerCamelCase    ,		)
        return block_records


 def 						a__   ( self						):
        __a       =    RealmRetriever(
            block_records=self.get_dummy_block_records()    ,		tokenizer=self.get_tokenizer()    ,		)
        return retriever


 def 						a__   ( self						):
        __a       =    self.get_config()
        __a       =    self.get_dummy_retriever()
        __a       =    retriever.tokenizer

        __a       =    np.array([0, 3]    ,		dtype="long"						)
        __a       =    tokenizer(["Test question"]						).input_ids
        __a       =    tokenizer(
            ["the fourth"]    ,		add_special_tokens=lowerCamelCase    ,		return_token_type_ids=lowerCamelCase    ,		return_attention_mask=lowerCamelCase    ,		).input_ids
        __a       =    config.reader_seq_len

        __a ,   __a ,   __a ,   __a       =    retriever(
            lowerCamelCase    ,		lowerCamelCase    ,		answer_ids=lowerCamelCase    ,		max_length=lowerCamelCase    ,		return_tensors="np"						)

        self.assertEqual(len(lowerCamelCase						)    ,		2						)
        self.assertEqual(len(lowerCamelCase						)    ,		2						)
        self.assertEqual(len(lowerCamelCase						)    ,		2						)
        self.assertEqual(concat_inputs.input_ids.shape    ,		(2, 10)						)
        self.assertEqual(concat_inputs.attention_mask.shape    ,		(2, 10)						)
        self.assertEqual(concat_inputs.token_type_ids.shape    ,		(2, 10)						)
        self.assertEqual(concat_inputs.special_tokens_mask.shape    ,		(2, 10)						)
        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						):
        __a       =    self.get_config()
        __a       =    self.get_dummy_retriever()
        __a       =    retriever.tokenizer

        __a       =    np.array([0, 3, 5]    ,		dtype="long"						)
        __a       =    tokenizer(["Test question"]						).input_ids
        __a       =    tokenizer(
            ["the fourth", "longer longer"]    ,		add_special_tokens=lowerCamelCase    ,		return_token_type_ids=lowerCamelCase    ,		return_attention_mask=lowerCamelCase    ,		).input_ids
        __a       =    config.reader_seq_len

        __a ,   __a ,   __a ,   __a       =    retriever(
            lowerCamelCase    ,		lowerCamelCase    ,		answer_ids=lowerCamelCase    ,		max_length=lowerCamelCase    ,		return_tensors="np"						)

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






 def 						a__   ( self						):
        __a       =    self.get_dummy_retriever()
        retriever.save_pretrained(os.path.join(self.tmpdirname    ,		"realm_block_records"						)						)

        # Test local path
        __a       =    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:
               __a       =    os.path.join(
                   os.path.join(self.tmpdirname    ,		"realm_block_records"						)    ,		_REALM_BLOCK_RECORDS_FILENAME						)
               __a       =    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 diffusers.pipelines.pipeline_utils import is_safetensors_compatible






class  snake_case__      (       unittest.TestCase       ):


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
            "unet/diffusion_pytorch_model.bin",
            # Removed: 'unet/diffusion_pytorch_model.safetensors',
        ]
        self.assertFalse(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
        ]
        self.assertTrue(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.bin",
            "safety_checker/model.safetensors",
            "vae/diffusion_pytorch_model.bin",
            "vae/diffusion_pytorch_model.safetensors",
            "text_encoder/pytorch_model.bin",
            # Removed: 'text_encoder/model.safetensors',
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        self.assertFalse(is_safetensors_compatible(lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        # pass variant but use the non-variant filenames
        __a       =    [
            "unet/diffusion_pytorch_model.bin",
            "unet/diffusion_pytorch_model.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
            "unet/diffusion_pytorch_model.fp16.bin",
            # Removed: 'unet/diffusion_pytorch_model.fp16.safetensors',
        ]
        __a       =    "fp16"
        self.assertFalse(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        __a       =    [
            "text_encoder/pytorch_model.fp16.bin",
            "text_encoder/model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)


 def 						a__   ( self						):
        # pass variant but use the non-variant filenames
        __a       =    [
            "text_encoder/pytorch_model.bin",
            "text_encoder/model.safetensors",
        ]
        __a       =    "fp16"
        self.assertTrue(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)






 def 						a__   ( self						):
        __a       =    [
            "safety_checker/pytorch_model.fp16.bin",
            "safety_checker/model.fp16.safetensors",
            "vae/diffusion_pytorch_model.fp16.bin",
            "vae/diffusion_pytorch_model.fp16.safetensors",
            "text_encoder/pytorch_model.fp16.bin",
            # 'text_encoder/model.fp16.safetensors',
            "unet/diffusion_pytorch_model.fp16.bin",
            "unet/diffusion_pytorch_model.fp16.safetensors",
        ]
        __a       =    "fp16"
        self.assertFalse(is_safetensors_compatible(lowerCamelCase    ,		variant=lowerCamelCase						)						)
