

"""simple docstring"""





import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Image
from datasets.packaged_modules.text.text import Text

from ..utils import require_pil




@pytest.fixture
def 						_A						(						UpperCamelCase_  : Dict)			->		List[str]:
   '''simple docstring'''


   __lowercase			=		tmp_path / "text.txt"
   __lowercase			=		textwrap.dedent(
       "\\n        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.\n        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.\n        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.\n        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.\n\n        Second paragraph:\n        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.\n        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.\n        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.\n        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.\n        ")
   with open(UpperCamelCase_,    "w",    encoding="utf-8") as f:
      f.write(UpperCamelCase_)
   return str(UpperCamelCase_)




@pytest.fixture
def 						_A						(						UpperCamelCase_  : Optional[Any],    UpperCamelCase_  : List[str])			->		Optional[Any]:
   '''simple docstring'''


   __lowercase			=		tmp_path / "text_with_image.txt"
   with open(UpperCamelCase_,    "w",    encoding="utf-8") as f:
      f.write(UpperCamelCase_)
   return str(UpperCamelCase_)




@pytest.mark.parametrize("keep_linebreaks",    [True, False])
def 						_A						(						UpperCamelCase_  : Union[str, Any],    UpperCamelCase_  : Optional[int])			->		Tuple:
   '''simple docstring'''


   with open(UpperCamelCase_,    encoding="utf-8") as f:
      __lowercase			=		f.read().splitlines(keepends=UpperCamelCase_)
   __lowercase			=		Text(keep_linebreaks=UpperCamelCase_,    encoding="utf-8")
   __lowercase			=		text._generate_tables([[text_file]])
   __lowercase			=		pa.concat_tables([table for _, table in generator]).to_pydict()["text"]
   assert generated_content == expected_content




@require_pil
def 						_A						(						UpperCamelCase_  : Tuple)			->		Union[str, Any]:
   '''simple docstring'''


   with open(UpperCamelCase_,    encoding="utf-8") as f:
      __lowercase			=		f.read().splitlines()[0]
   __lowercase			=		Text(encoding="utf-8",    features=Features({"image": Image()}))
   __lowercase			=		text._generate_tables([[text_file_with_image]])
   __lowercase			=		pa.concat_tables([table for _, table in generator])
   assert pa_table.schema.field("image").type == Image()()
   __lowercase			=		pa_table.to_pydict()["image"]
   assert generated_content == [{"path": image_file, "bytes": None}]




@pytest.mark.parametrize("sample_by",    ["line", "paragraph", "document"])
def 						_A						(						UpperCamelCase_  : Any,    UpperCamelCase_  : Dict)			->		Tuple:
   '''simple docstring'''


   with open(UpperCamelCase_,    encoding="utf-8") as f:
      __lowercase			=		f.read()
   if sample_by == "line":
      __lowercase			=		expected_content.splitlines()
   elif sample_by == "paragraph":
      __lowercase			=		expected_content.split("\n\n")
   elif sample_by == "document":
      __lowercase			=		[expected_content]
   __lowercase			=		Text(sample_by=UpperCamelCase_,    encoding="utf-8",    chunksize=100)
   __lowercase			=		text._generate_tables([[text_file]])
   __lowercase			=		pa.concat_tables([table for _, table in generator]).to_pydict()["text"]
   assert generated_content == expected_content







"""simple docstring"""





import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp




def 						_A						(						UpperCamelCase_  : int,    UpperCamelCase_  : Union[str, Any],    UpperCamelCase_  : List[Any],    UpperCamelCase_  : Dict,    UpperCamelCase_  : int = 4096)			->		Optional[Any]:
   '''simple docstring'''


   __lowercase	,__lowercase	,__lowercase			=		key.shape[-3:]
   __lowercase			=		value.shape[-1]
   __lowercase			=		min(UpperCamelCase_,    UpperCamelCase_)
   __lowercase			=		query / jnp.sqrt(UpperCamelCase_)

   @functools.partial(jax.checkpoint,    prevent_cse=UpperCamelCase_)
   def summarize_chunk(UpperCamelCase_  : int,    UpperCamelCase_  : Union[str, Any],    UpperCamelCase_  : Any):
      __lowercase			=		jnp.einsum("...qhd,...khd->...qhk",    UpperCamelCase_,    UpperCamelCase_,    precision=UpperCamelCase_)

      __lowercase			=		jnp.max(UpperCamelCase_,    axis=-1,    keepdims=UpperCamelCase_)
      __lowercase			=		jax.lax.stop_gradient(UpperCamelCase_)
      __lowercase			=		jnp.exp(attn_weights - max_score)

      __lowercase			=		jnp.einsum("...vhf,...qhv->...qhf",    UpperCamelCase_,    UpperCamelCase_,    precision=UpperCamelCase_)
      __lowercase			=		jnp.einsum("...qhk->...qh",    UpperCamelCase_)

      return (exp_values, exp_weights.sum(axis=-1), max_score)

   def chunk_scanner(UpperCamelCase_  : int):
      # julienne key array
      __lowercase			=		jax.lax.dynamic_slice(
          operand=UpperCamelCase_,    start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0],    slice_sizes=list(key.shape[:-3]) + [key_chunk_size, num_heads, k_features],    )

      # julienne value array
      __lowercase			=		jax.lax.dynamic_slice(
          operand=UpperCamelCase_,    start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0],    slice_sizes=list(value.shape[:-3]) + [key_chunk_size, num_heads, v_features],    )

      return summarize_chunk(UpperCamelCase_,    UpperCamelCase_,    UpperCamelCase_)

   __lowercase	,__lowercase	,__lowercase			=		jax.lax.map(f=UpperCamelCase_,    xs=jnp.arange(0,    UpperCamelCase_,    UpperCamelCase_))

   __lowercase			=		jnp.max(UpperCamelCase_,    axis=0,    keepdims=UpperCamelCase_)
   __lowercase			=		jnp.exp(chunk_max - global_max)

   chunk_values *= jnp.expand_dims(UpperCamelCase_,    axis=-1)
   chunk_weights *= max_diffs

   __lowercase			=		chunk_values.sum(axis=0)
   __lowercase			=		jnp.expand_dims(UpperCamelCase_,    -1).sum(axis=0)

   return all_values / all_weights




def 						_A						(						UpperCamelCase_  : Optional[Any],    UpperCamelCase_  : Optional[int],    UpperCamelCase_  : List[str],    UpperCamelCase_  : Union[str, Any]=jax.lax.Precision.HIGHEST,    UpperCamelCase_  : int = 1024,    UpperCamelCase_  : int = 4096)			->		str:
   '''simple docstring'''


   __lowercase	,__lowercase	,__lowercase			=		query.shape[-3:]

   def chunk_scanner(UpperCamelCase_  : List[str],    UpperCamelCase_  : str):
      # julienne query array
      __lowercase			=		jax.lax.dynamic_slice(
          operand=UpperCamelCase_,    start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0],    slice_sizes=list(query.shape[:-3]) + [min(UpperCamelCase_,    UpperCamelCase_), num_heads, q_features],    )

      return (
          chunk_idx + query_chunk_size,  # unused ignore it
          _query_chunk_attention(
              query=UpperCamelCase_,    key=UpperCamelCase_,    value=UpperCamelCase_,    precision=UpperCamelCase_,    key_chunk_size=UpperCamelCase_),
      )

   __lowercase	,__lowercase			=		jax.lax.scan(
       f=UpperCamelCase_,    init=0,    xs=UpperCamelCase_,    length=math.ceil(num_q / query_chunk_size))

   return jnp.concatenate(UpperCamelCase_,    axis=-3)  # fuse the chunked result back




class       _lowerCAmelCase						(			nn.Module     ):







 """simple docstring"""







 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      int     	=			8
 __UpperCAmelCase			:      int     	=			6_4
 __UpperCAmelCase			:      float     	=			0.0
 __UpperCAmelCase			:      bool     	=			False
 __UpperCAmelCase			:      jnp.dtype     	=			jnp.floataa


 def       _lowercase				(    self							:    List[Any]			):
    __lowercase			=		self.dim_head * self.heads
    __lowercase			=		self.dim_head**-0.5

    # Weights were exported with old names {to_q, to_k, to_v, to_out}
    __lowercase			=		nn.Dense(UpperCAmelCase__,      use_bias=UpperCAmelCase__,      dtype=self.dtype,      name="to_q"			)
    __lowercase			=		nn.Dense(UpperCAmelCase__,      use_bias=UpperCAmelCase__,      dtype=self.dtype,      name="to_k"			)
    __lowercase			=		nn.Dense(UpperCAmelCase__,      use_bias=UpperCAmelCase__,      dtype=self.dtype,      name="to_v"			)

    __lowercase			=		nn.Dense(self.query_dim,      dtype=self.dtype,      name="to_out_0"			)
    __lowercase			=		nn.Dropout(rate=self.dropout			)


 def       _lowercase				(    self							:    Any,      UpperCAmelCase__							:    Union[str, Any]			):
    __lowercase	,__lowercase	,__lowercase			=		tensor.shape
    __lowercase			=		self.heads
    __lowercase			=		tensor.reshape(UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__,      dim // head_size			)
    __lowercase			=		jnp.transpose(UpperCAmelCase__,      (0, 2, 1, 3)			)
    __lowercase			=		tensor.reshape(batch_size * head_size,      UpperCAmelCase__,      dim // head_size			)
    return tensor


 def       _lowercase				(    self							:    Union[str, Any],      UpperCAmelCase__							:    int			):
    __lowercase	,__lowercase	,__lowercase			=		tensor.shape
    __lowercase			=		self.heads
    __lowercase			=		tensor.reshape(batch_size // head_size,      UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__			)
    __lowercase			=		jnp.transpose(UpperCAmelCase__,      (0, 2, 1, 3)			)
    __lowercase			=		tensor.reshape(batch_size // head_size,      UpperCAmelCase__,      dim * head_size			)
    return tensor



 def __call__(    self							:    List[Any],      UpperCAmelCase__							:    str,      UpperCAmelCase__							:    Union[str, Any]=None,      UpperCAmelCase__							:    Optional[Any]=True			):
    __lowercase			=		hidden_states if context is None else context

    __lowercase			=		self.query(UpperCAmelCase__			)
    __lowercase			=		self.key(UpperCAmelCase__			)
    __lowercase			=		self.value(UpperCAmelCase__			)

    __lowercase			=		self.reshape_heads_to_batch_dim(UpperCAmelCase__			)
    __lowercase			=		self.reshape_heads_to_batch_dim(UpperCAmelCase__			)
    __lowercase			=		self.reshape_heads_to_batch_dim(UpperCAmelCase__			)

    if self.use_memory_efficient_attention:
       __lowercase			=		query_states.transpose(1,      0,      2			)
       __lowercase			=		key_states.transpose(1,      0,      2			)
       __lowercase			=		value_states.transpose(1,      0,      2			)

       # this if statement create a chunk size for each layer of the unet
       # the chunk size is equal to the query_length dimension of the deepest layer of the unet

       __lowercase			=		query_states.shape[-3]
       if flatten_latent_dim % 6_4 == 0:
          __lowercase			=		int(flatten_latent_dim / 6_4			)
       elif flatten_latent_dim % 1_6 == 0:
          __lowercase			=		int(flatten_latent_dim / 1_6			)
       elif flatten_latent_dim % 4 == 0:
          __lowercase			=		int(flatten_latent_dim / 4			)
       else:
          __lowercase			=		int(UpperCAmelCase__			)

       __lowercase			=		jax_memory_efficient_attention(
           UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__,      query_chunk_size=UpperCAmelCase__,      key_chunk_size=4_0_9_6 * 4			)

       __lowercase			=		hidden_states.transpose(1,      0,      2			)
    else:
       # compute attentions
       __lowercase			=		jnp.einsum("b i d, b j d->b i j",      UpperCAmelCase__,      UpperCAmelCase__			)
       __lowercase			=		attention_scores * self.scale
       __lowercase			=		nn.softmax(UpperCAmelCase__,      axis=2			)

       # attend to values
       __lowercase			=		jnp.einsum("b i j, b j d -> b i d",      UpperCAmelCase__,      UpperCAmelCase__			)

    __lowercase			=		self.reshape_batch_dim_to_heads(UpperCAmelCase__			)
    __lowercase			=		self.proj_attn(UpperCAmelCase__			)
    return self.dropout_layer(UpperCAmelCase__,      deterministic=UpperCAmelCase__			)







class       _lowerCAmelCase						(			nn.Module     ):







 """simple docstring"""







 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      float     	=			0.0
 __UpperCAmelCase			:      bool     	=			False
 __UpperCAmelCase			:      jnp.dtype     	=			jnp.floataa
 __UpperCAmelCase			:      bool     	=			False


 def       _lowercase				(    self							:    Tuple			):
    # self attention (or cross_attention if only_cross_attention is True)
    __lowercase			=		FlaxAttention(
        self.dim,      self.n_heads,      self.d_head,      self.dropout,      self.use_memory_efficient_attention,      dtype=self.dtype			)
    # cross attention
    __lowercase			=		FlaxAttention(
        self.dim,      self.n_heads,      self.d_head,      self.dropout,      self.use_memory_efficient_attention,      dtype=self.dtype			)
    __lowercase			=		FlaxFeedForward(dim=self.dim,      dropout=self.dropout,      dtype=self.dtype			)
    __lowercase			=		nn.LayerNorm(epsilon=1E-5,      dtype=self.dtype			)
    __lowercase			=		nn.LayerNorm(epsilon=1E-5,      dtype=self.dtype			)
    __lowercase			=		nn.LayerNorm(epsilon=1E-5,      dtype=self.dtype			)
    __lowercase			=		nn.Dropout(rate=self.dropout			)



 def __call__(    self							:    str,      UpperCAmelCase__							:    Union[str, Any],      UpperCAmelCase__							:    Optional[Any],      UpperCAmelCase__							:    Dict=True			):
    # self attention
    __lowercase			=		hidden_states
    if self.only_cross_attention:
       __lowercase			=		self.attna(self.norma(UpperCAmelCase__			),      UpperCAmelCase__,      deterministic=UpperCAmelCase__			)
    else:
       __lowercase			=		self.attna(self.norma(UpperCAmelCase__			),      deterministic=UpperCAmelCase__			)
    __lowercase			=		hidden_states + residual

    # cross attention
    __lowercase			=		hidden_states
    __lowercase			=		self.attna(self.norma(UpperCAmelCase__			),      UpperCAmelCase__,      deterministic=UpperCAmelCase__			)
    __lowercase			=		hidden_states + residual

    # feed forward
    __lowercase			=		hidden_states
    __lowercase			=		self.ff(self.norma(UpperCAmelCase__			),      deterministic=UpperCAmelCase__			)
    __lowercase			=		hidden_states + residual

    return self.dropout_layer(UpperCAmelCase__,      deterministic=UpperCAmelCase__			)







class       _lowerCAmelCase						(			nn.Module     ):







 """simple docstring"""







 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      int     	=			1
 __UpperCAmelCase			:      float     	=			0.0
 __UpperCAmelCase			:      bool     	=			False
 __UpperCAmelCase			:      bool     	=			False
 __UpperCAmelCase			:      jnp.dtype     	=			jnp.floataa
 __UpperCAmelCase			:      bool     	=			False


 def       _lowercase				(    self							:    Dict			):
    __lowercase			=		nn.GroupNorm(num_groups=3_2,      epsilon=1E-5			)

    __lowercase			=		self.n_heads * self.d_head
    if self.use_linear_projection:
       __lowercase			=		nn.Dense(UpperCAmelCase__,      dtype=self.dtype			)
    else:
       __lowercase			=		nn.Conv(
           UpperCAmelCase__,      kernel_size=(1, 1),      strides=(1, 1),      padding="VALID",      dtype=self.dtype,      )

    __lowercase			=		[
        FlaxBasicTransformerBlock(
            UpperCAmelCase__,      self.n_heads,      self.d_head,      dropout=self.dropout,      only_cross_attention=self.only_cross_attention,      dtype=self.dtype,      use_memory_efficient_attention=self.use_memory_efficient_attention,      )
        for _ in range(self.depth			)
    ]

    if self.use_linear_projection:
       __lowercase			=		nn.Dense(UpperCAmelCase__,      dtype=self.dtype			)
    else:
       __lowercase			=		nn.Conv(
           UpperCAmelCase__,      kernel_size=(1, 1),      strides=(1, 1),      padding="VALID",      dtype=self.dtype,      )

    __lowercase			=		nn.Dropout(rate=self.dropout			)



 def __call__(    self							:    Dict,      UpperCAmelCase__							:    Union[str, Any],      UpperCAmelCase__							:    Tuple,      UpperCAmelCase__							:    List[str]=True			):
    __lowercase	,__lowercase	,__lowercase	,__lowercase			=		hidden_states.shape
    __lowercase			=		hidden_states
    __lowercase			=		self.norm(UpperCAmelCase__			)
    if self.use_linear_projection:
       __lowercase			=		hidden_states.reshape(UpperCAmelCase__,      height * width,      UpperCAmelCase__			)
       __lowercase			=		self.proj_in(UpperCAmelCase__			)
    else:
       __lowercase			=		self.proj_in(UpperCAmelCase__			)
       __lowercase			=		hidden_states.reshape(UpperCAmelCase__,      height * width,      UpperCAmelCase__			)

    for transformer_block in self.transformer_blocks:
       __lowercase			=		transformer_block(UpperCAmelCase__,      UpperCAmelCase__,      deterministic=UpperCAmelCase__			)

    if self.use_linear_projection:
       __lowercase			=		self.proj_out(UpperCAmelCase__			)
       __lowercase			=		hidden_states.reshape(UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__			)
    else:
       __lowercase			=		hidden_states.reshape(UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__,      UpperCAmelCase__			)
       __lowercase			=		self.proj_out(UpperCAmelCase__			)

    __lowercase			=		hidden_states + residual
    return self.dropout_layer(UpperCAmelCase__,      deterministic=UpperCAmelCase__			)







class       _lowerCAmelCase						(			nn.Module     ):







 """simple docstring"""







 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      float     	=			0.0
 __UpperCAmelCase			:      jnp.dtype     	=			jnp.floataa


 def       _lowercase				(    self							:    List[str]			):
    # The second linear layer needs to be called
    # net_2 for now to match the index of the Sequential layer
    __lowercase			=		FlaxGEGLU(self.dim,      self.dropout,      self.dtype			)
    __lowercase			=		nn.Dense(self.dim,      dtype=self.dtype			)



 def __call__(    self							:    int,      UpperCAmelCase__							:    List[str],      UpperCAmelCase__							:    Union[str, Any]=True			):
    __lowercase			=		self.net_a(UpperCAmelCase__,      deterministic=UpperCAmelCase__			)
    __lowercase			=		self.net_a(UpperCAmelCase__			)
    return hidden_states







class       _lowerCAmelCase						(			nn.Module     ):







 """simple docstring"""







 __UpperCAmelCase			:      int
 __UpperCAmelCase			:      float     	=			0.0
 __UpperCAmelCase			:      jnp.dtype     	=			jnp.floataa


 def       _lowercase				(    self							:    int			):
    __lowercase			=		self.dim * 4
    __lowercase			=		nn.Dense(inner_dim * 2,      dtype=self.dtype			)
    __lowercase			=		nn.Dropout(rate=self.dropout			)



 def __call__(    self							:    Optional[Any],      UpperCAmelCase__							:    str,      UpperCAmelCase__							:    int=True			):
    __lowercase			=		self.proj(UpperCAmelCase__			)
    __lowercase	,__lowercase			=		jnp.split(UpperCAmelCase__,      2,      axis=2			)
    return self.dropout_layer(hidden_linear * nn.gelu(UpperCAmelCase__			),      deterministic=UpperCAmelCase__			)




