





"""simple docstring"""


from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings



@maybe_allow_in_graph
class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       Optional[Any]					,   __a   :       int					,   __a   :       int					,   __a   :       int					,   __a   :       Dict=0.0					,   __a   :       Optional[int] = None					,   __a   :       str = "geglu"					,   __a   :       Optional[int] = None					,   __a   :       bool = False					,   __a   :       bool = False					,   __a   :       bool = False					,   __a   :       bool = False					,   __a   :       bool = True					,   __a   :       str = "layer_norm"					,   __a   :       bool = False					,   )		->				List[str]:
        super().__init__()
        _UpperCamelCase      :	Tuple						=       only_cross_attention

        _UpperCamelCase      :	List[Any]						=       (num_embeds_ada_norm is not None) and norm_type == "ada_norm_zero"
        _UpperCamelCase      :	List[Any]						=       (num_embeds_ada_norm is not None) and norm_type == "ada_norm"

        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
              raise ValueError(
                  F'''`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'''
                  F''' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'''	)

        # Define 3 blocks. Each block has its own normalization layer.
        # 1. Self-Attn
        if self.use_ada_layer_norm:
              _UpperCamelCase      :	Union[str, Any]						=       AdaLayerNorm(__a					,   __a	)
        elif self.use_ada_layer_norm_zero:
              _UpperCamelCase      :	Any						=       AdaLayerNormZero(__a					,   __a	)
        else:
              _UpperCamelCase      :	Optional[Any]						=       nn.LayerNorm(__a					,   elementwise_affine=__a	)
        _UpperCamelCase      :	Union[str, Any]						=       Attention(
            query_dim=__a					,   heads=__a					,   dim_head=__a					,   dropout=__a					,   bias=__a					,   cross_attention_dim=cross_attention_dim if only_cross_attention else None					,   upcast_attention=__a					,   )

        # 2. Cross-Attn
        if cross_attention_dim is not None or double_self_attention:
              # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
              # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
              # the second cross attention block.
              _UpperCamelCase      :	int						=       (
                  AdaLayerNorm(__a					,   __a	)
                  if self.use_ada_layer_norm
                  else nn.LayerNorm(__a					,   elementwise_affine=__a	)
              )
              _UpperCamelCase      :	Optional[Any]						=       Attention(
                  query_dim=__a					,   cross_attention_dim=cross_attention_dim if not double_self_attention else None					,   heads=__a					,   dim_head=__a					,   dropout=__a					,   bias=__a					,   upcast_attention=__a					,   )  # is self-attn if encoder_hidden_states is none
        else:
              _UpperCamelCase      :	Optional[Any]						=       None
              _UpperCamelCase      :	Optional[int]						=       None

        # 3. Feed-forward
        _UpperCamelCase      :	List[Any]						=       nn.LayerNorm(__a					,   elementwise_affine=__a	)
        _UpperCamelCase      :	Any						=       FeedForward(__a					,   dropout=__a					,   activation_fn=__a					,   final_dropout=__a	)

        # let chunk size default to None
        _UpperCamelCase      :	List[str]						=       None
        _UpperCamelCase      :	Optional[Any]						=       0

  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]					,   __a   :       Optional[int]					,   __a   :       int	)		->				Tuple:
        # Sets chunk feed-forward
        _UpperCamelCase      :	List[str]						=       chunk_size
        _UpperCamelCase      :	Tuple						=       dim


  def    __SCREAMING_SNAKE_CASE				(		self   :       List[Any]					,   __a   :       torch.FloatTensor					,   __a   :       Optional[torch.FloatTensor] = None					,   __a   :       Optional[torch.FloatTensor] = None					,   __a   :       Optional[torch.FloatTensor] = None					,   __a   :       Optional[torch.LongTensor] = None					,   __a   :       Dict[str, Any] = None					,   __a   :       Optional[torch.LongTensor] = None					,   )		->				Optional[Any]:
        # Notice that normalization is always applied before the real computation in the following blocks.
        # 1. Self-Attention
        if self.use_ada_layer_norm:
              _UpperCamelCase      :	int						=       self.norma(__a					,   __a	)
        elif self.use_ada_layer_norm_zero:
              _UpperCamelCase      :	Dict						=       self.norma(
                  __a					,   __a					,   __a					,   hidden_dtype=hidden_states.dtype	)
        else:
              _UpperCamelCase      :	str						=       self.norma(__a	)

        _UpperCamelCase      :	Optional[int]						=       cross_attention_kwargs if cross_attention_kwargs is not None else {}

        _UpperCamelCase      :	Optional[int]						=       self.attna(
            __a					,   encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None					,   attention_mask=__a					,   **__a					,   )
        if self.use_ada_layer_norm_zero:
              _UpperCamelCase      :	Optional[int]						=       gate_msa.unsqueeze(1	) * attn_output
        _UpperCamelCase      :	Tuple						=       attn_output + hidden_states

        # 2. Cross-Attention
        if self.attna is not None:
              _UpperCamelCase      :	str						=       (
                  self.norma(__a					,   __a	) if self.use_ada_layer_norm else self.norma(__a	)
              )

              _UpperCamelCase      :	Any						=       self.attna(
                  __a					,   encoder_hidden_states=__a					,   attention_mask=__a					,   **__a					,   )
              _UpperCamelCase      :	List[Any]						=       attn_output + hidden_states

        # 3. Feed-forward
        _UpperCamelCase      :	Optional[int]						=       self.norma(__a	)

        if self.use_ada_layer_norm_zero:
              _UpperCamelCase      :	str						=       norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

        if self._chunk_size is not None:
              # "feed_forward_chunk_size" can be used to save memory
              if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
                    raise ValueError(
                        F'''`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'''	)

              _UpperCamelCase      :	str						=       norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
              _UpperCamelCase      :	Dict						=       torch.cat(
                  [self.ff(__a	) for hid_slice in norm_hidden_states.chunk(__a					,   dim=self._chunk_dim	)]					,   dim=self._chunk_dim					,   )
        else:
              _UpperCamelCase      :	Optional[int]						=       self.ff(__a	)

        if self.use_ada_layer_norm_zero:
              _UpperCamelCase      :	Any						=       gate_mlp.unsqueeze(1	) * ff_output

        _UpperCamelCase      :	Tuple						=       ff_output + hidden_states

        return hidden_states





class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       Optional[int]					,   __a   :       int					,   __a   :       Optional[int] = None					,   __a   :       int = 4					,   __a   :       float = 0.0					,   __a   :       str = "geglu"					,   __a   :       bool = False					,   )		->				str:
        super().__init__()
        _UpperCamelCase      :	List[Any]						=       int(dim * mult	)
        _UpperCamelCase      :	List[Any]						=       dim_out if dim_out is not None else dim

        if activation_fn == "gelu":
              _UpperCamelCase      :	Optional[Any]						=       GELU(__a					,   __a	)
        if activation_fn == "gelu-approximate":
              _UpperCamelCase      :	Dict						=       GELU(__a					,   __a					,   approximate="tanh"	)
        elif activation_fn == "geglu":
              _UpperCamelCase      :	Dict						=       GEGLU(__a					,   __a	)
        elif activation_fn == "geglu-approximate":
              _UpperCamelCase      :	List[str]						=       ApproximateGELU(__a					,   __a	)

        _UpperCamelCase      :	Union[str, Any]						=       nn.ModuleList([]	)
        # project in
        self.net.append(__a	)
        # project dropout
        self.net.append(nn.Dropout(__a	)	)
        # project out
        self.net.append(nn.Linear(__a					,   __a	)	)
        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
        if final_dropout:
              self.net.append(nn.Dropout(__a	)	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple					,   __a   :       Tuple	)		->				List[str]:
        for module in self.net:
              _UpperCamelCase      :	int						=       module(__a	)
        return hidden_states





class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       int					,   __a   :       int					,   __a   :       int					,   __a   :       str = "none"	)		->				List[str]:
        super().__init__()
        _UpperCamelCase      :	Optional[Any]						=       nn.Linear(__a					,   __a	)
        _UpperCamelCase      :	List[str]						=       approximate

  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[int]					,   __a   :       int	)		->				Dict:
        if gate.device.type != "mps":
              return F.gelu(__a					,   approximate=self.approximate	)
        # mps: gelu is not implemented for float16
        return F.gelu(gate.to(dtype=torch.floataa	)					,   approximate=self.approximate	).to(dtype=gate.dtype	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]					,   __a   :       int	)		->				Tuple:
        _UpperCamelCase      :	int						=       self.proj(__a	)
        _UpperCamelCase      :	List[Any]						=       self.gelu(__a	)
        return hidden_states





class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       List[str]					,   __a   :       int					,   __a   :       int	)		->				str:
        super().__init__()
        _UpperCamelCase      :	Optional[Any]						=       nn.Linear(__a					,   dim_out * 2	)

  def    __SCREAMING_SNAKE_CASE				(		self   :       Any					,   __a   :       Optional[Any]	)		->				List[str]:
        if gate.device.type != "mps":
              return F.gelu(__a	)
        # mps: gelu is not implemented for float16
        return F.gelu(gate.to(dtype=torch.floataa	)	).to(dtype=gate.dtype	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       Tuple					,   __a   :       Tuple	)		->				Optional[int]:
        _UpperCamelCase      :	Optional[Any]						=       self.proj(__a	).chunk(2					,   dim=-1	)
        return hidden_states * self.gelu(__a	)





class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       int					,   __a   :       int					,   __a   :       int	)		->				Any:
        super().__init__()
        _UpperCamelCase      :	str						=       nn.Linear(__a					,   __a	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       Union[str, Any]					,   __a   :       Any	)		->				str:
        _UpperCamelCase      :	Dict						=       self.proj(__a	)
        return x * torch.sigmoid(1.7_02 * x	)





class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       Dict					,   __a   :       List[str]					,   __a   :       Tuple	)		->				Union[str, Any]:
        super().__init__()
        _UpperCamelCase      :	Dict						=       nn.Embedding(__a					,   __a	)
        _UpperCamelCase      :	Optional[int]						=       nn.SiLU()
        _UpperCamelCase      :	Tuple						=       nn.Linear(__a					,   embedding_dim * 2	)
        _UpperCamelCase      :	str						=       nn.LayerNorm(__a					,   elementwise_affine=__a	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       Any					,   __a   :       str					,   __a   :       List[Any]	)		->				List[Any]:
        _UpperCamelCase      :	Optional[int]						=       self.linear(self.silu(self.emb(__a	)	)	)
        _UpperCamelCase      :	Optional[Any]						=       torch.chunk(__a					,   2	)
        _UpperCamelCase      :	List[Any]						=       self.norm(__a	) * (1 + scale) + shift
        return x





class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       Tuple					,   __a   :       List[Any]					,   __a   :       Tuple	)		->				List[Any]:
        super().__init__()

        _UpperCamelCase      :	Any						=       CombinedTimestepLabelEmbeddings(__a					,   __a	)

        _UpperCamelCase      :	List[str]						=       nn.SiLU()
        _UpperCamelCase      :	Optional[int]						=       nn.Linear(__a					,   6 * embedding_dim					,   bias=__a	)
        _UpperCamelCase      :	List[str]						=       nn.LayerNorm(__a					,   elementwise_affine=__a					,   eps=1e-6	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       Optional[int]					,   __a   :       str					,   __a   :       Optional[Any]					,   __a   :       Dict					,   __a   :       Tuple=None	)		->				int:
        _UpperCamelCase      :	Tuple						=       self.linear(self.silu(self.emb(__a					,   __a					,   hidden_dtype=__a	)	)	)
        _UpperCamelCase      :	str						=       emb.chunk(6					,   dim=1	)
        _UpperCamelCase      :	Tuple						=       self.norm(__a	) * (1 + scale_msa[:, None]) + shift_msa[:, None]
        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp





class      __SCREAMING_SNAKE_CASE    (						nn.Module   ):







  '''simple docstring'''


  def __init__(		self   :       str					,   __a   :       int					,   __a   :       int					,   __a   :       int					,   __a   :       Optional[str] = None					,   __a   :       float = 1e-5	)		->				Dict:
        super().__init__()
        _UpperCamelCase      :	Optional[int]						=       num_groups
        _UpperCamelCase      :	List[str]						=       eps

        if act_fn is None:
              _UpperCamelCase      :	int						=       None
        else:
              _UpperCamelCase      :	List[str]						=       get_activation(__a	)

        _UpperCamelCase      :	List[Any]						=       nn.Linear(__a					,   out_dim * 2	)


  def    __SCREAMING_SNAKE_CASE				(		self   :       List[str]					,   __a   :       Union[str, Any]					,   __a   :       Any	)		->				Tuple:
        if self.act:
              _UpperCamelCase      :	Any						=       self.act(__a	)
        _UpperCamelCase      :	Optional[Any]						=       self.linear(__a	)
        _UpperCamelCase      :	Tuple						=       emb[:, :, None, None]
        _UpperCamelCase      :	Optional[int]						=       emb.chunk(2					,   dim=1	)

        _UpperCamelCase      :	Optional[int]						=       F.group_norm(__a					,   self.num_groups					,   eps=self.eps	)
        _UpperCamelCase      :	Optional[Any]						=       x * (1 + scale) + shift
        return x







"""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,  _UpperCamelCase,  _UpperCamelCase,  _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,  _UpperCamelCase,  _UpperCamelCase,  _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"	)
