





"""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      __SCREAMING_SNAKE_CASE    (						_UpperCamelCase   ):







  '''simple docstring'''


  def    __SCREAMING_SNAKE_CASE				(		self   :       Dict	)		->				Optional[Any]:
        _UpperCamelCase      :	Tuple						=       tempfile.mkdtemp()
        _UpperCamelCase      :	str						=       5

        # Realm tok
        _UpperCamelCase      :	Tuple						=       [
            "[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      :	Optional[int]						=       os.path.join(self.tmpdirname					,   "realm_tokenizer"	)
        os.makedirs(__a					,   exist_ok=__a	)
        _UpperCamelCase      :	Optional[Any]						=       os.path.join(__a					,   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      :	Optional[int]						=       os.path.join(self.tmpdirname					,   "realm_block_records"	)
        os.makedirs(__a					,   exist_ok=__a	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       str	)		->				RealmTokenizer:
        return RealmTokenizer.from_pretrained(os.path.join(self.tmpdirname					,   "realm_tokenizer"	)	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]	)		->				Dict:
        shutil.rmtree(self.tmpdirname	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       Union[str, Any]	)		->				List[str]:
        _UpperCamelCase      :	Optional[Any]						=       RealmConfig(num_block_records=self.num_block_records	)
        return config

  def    __SCREAMING_SNAKE_CASE				(		self   :       int	)		->				int:
        _UpperCamelCase      :	Any						=       Dataset.from_dict(
            {
                "id": ["0", "1"],
                "question": ["foo", "bar"],
                "answers": [["Foo", "Bar"], ["Bar"]],
            }	)
        return dataset

  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]	)		->				str:
        _UpperCamelCase      :	int						=       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=__a					,   )
        return block_records

  def    __SCREAMING_SNAKE_CASE				(		self   :       List[Any]	)		->				Optional[Any]:
        _UpperCamelCase      :	List[str]						=       RealmRetriever(
            block_records=self.get_dummy_block_records()					,   tokenizer=self.get_tokenizer()					,   )
        return retriever

  def    __SCREAMING_SNAKE_CASE				(		self   :       List[Any]	)		->				Tuple:
        _UpperCamelCase      :	Tuple						=       self.get_config()
        _UpperCamelCase      :	int						=       self.get_dummy_retriever()
        _UpperCamelCase      :	Tuple						=       retriever.tokenizer

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

        _UpperCamelCase      :	List[Any]						=       retriever(
            __a					,   __a					,   answer_ids=__a					,   max_length=__a					,   return_tensors="np"	)

        self.assertEqual(len(__a	)					,   2	)
        self.assertEqual(len(__a	)					,   2	)
        self.assertEqual(len(__a	)					,   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    __SCREAMING_SNAKE_CASE				(		self   :       Union[str, Any]	)		->				List[Any]:
        _UpperCamelCase      :	Any						=       self.get_config()
        _UpperCamelCase      :	Dict						=       self.get_dummy_retriever()
        _UpperCamelCase      :	Dict						=       retriever.tokenizer

        _UpperCamelCase      :	List[Any]						=       np.array([0, 3, 5]					,   dtype="long"	)
        _UpperCamelCase      :	Optional[int]						=       tokenizer(["Test question"]	).input_ids
        _UpperCamelCase      :	str						=       tokenizer(
            ["the fourth", "longer longer"]					,   add_special_tokens=__a					,   return_token_type_ids=__a					,   return_attention_mask=__a					,   ).input_ids
        _UpperCamelCase      :	Union[str, Any]						=       config.reader_seq_len

        _UpperCamelCase      :	Optional[Any]						=       retriever(
            __a					,   __a					,   answer_ids=__a					,   max_length=__a					,   return_tensors="np"	)

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


  def    __SCREAMING_SNAKE_CASE				(		self   :       int	)		->				List[str]:
        _UpperCamelCase      :	List[Any]						=       self.get_dummy_retriever()
        retriever.save_pretrained(os.path.join(self.tmpdirname					,   "realm_block_records"	)	)

        # Test local path
        _UpperCamelCase      :	int						=       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[Any]						=       os.path.join(
                  os.path.join(self.tmpdirname					,   "realm_block_records"	)					,   _REALM_BLOCK_RECORDS_FILENAME	)
              _UpperCamelCase      :	int						=       RealmRetriever.from_pretrained("google/realm-cc-news-pretrained-openqa"	)

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












"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


lowerCamelCase__				  =   {
    "configuration_owlvit": [
        "OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "OwlViTConfig",
        "OwlViTOnnxConfig",
        "OwlViTTextConfig",
        "OwlViTVisionConfig",
    ],
    "processing_owlvit": ["OwlViTProcessor"],
}


try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__				  =   ["OwlViTFeatureExtractor"]
  lowerCamelCase__				  =   ["OwlViTImageProcessor"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowerCamelCase__				  =   [
      "OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "OwlViTModel",
      "OwlViTPreTrainedModel",
      "OwlViTTextModel",
      "OwlViTVisionModel",
      "OwlViTForObjectDetection",
  ]

if TYPE_CHECKING:
  from .configuration_owlvit import (
      OWLVIT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      OwlViTConfig,
      OwlViTOnnxConfig,
      OwlViTTextConfig,
      OwlViTVisionConfig,
  )
  from .processing_owlvit import OwlViTProcessor

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .feature_extraction_owlvit import OwlViTFeatureExtractor
    from .image_processing_owlvit import OwlViTImageProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_owlvit import (
        OWLVIT_PRETRAINED_MODEL_ARCHIVE_LIST,
        OwlViTForObjectDetection,
        OwlViTModel,
        OwlViTPreTrainedModel,
        OwlViTTextModel,
        OwlViTVisionModel,
    )

else:
  import sys

  lowerCamelCase__				  =   _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
