


from unittest.mock import Mock, patch

from file_transfer.send_file import send_file







@patch('socket.socket'				)
@patch('builtins.open'				)
def        SCREAMING_SNAKE_CASE_						(			__A :			int  ,    __A :			Any				)	->		Optional[int]:

       """simple docstring"""

       a_					:		Any							=  Mock()
       a_					:		Dict							=  conn, Mock()
       a_					:		Optional[int]							=  iter([1, None]				)
       a_					:		List[str]							=  lambda __A				: next(__A				)

       # ===== invoke =====
       send_file(filename='mytext.txt'  ,    testing=__A				)

       # ===== ensurance =====
       sock.assert_called_once()
       sock.return_value.bind.assert_called_once()
       sock.return_value.listen.assert_called_once()
       sock.return_value.accept.assert_called_once()
       conn.recv.assert_called_once()

       file.return_value.__enter__.assert_called_once()
       file.return_value.__enter__.return_value.read.assert_called()

       conn.send.assert_called_once()
       conn.close.assert_called_once()
       sock.return_value.shutdown.assert_called_once()
       sock.return_value.close.assert_called_once()









import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
     import os

     # The slow tests are often failing with OOM error on GPU
     # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
     # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
     __A									=				"platform"

     import jax
     import jax.numpy as jnp

     from transformers import BlenderbotTokenizer
     from transformers.models.blenderbot.modeling_flax_blenderbot import (
         FlaxBlenderbotForConditionalGeneration,
         FlaxBlenderbotModel,
         shift_tokens_right,
     )





def      lowerCAmelCase_ (       __a	,    __a	,    __a=None	,    __a=None	,    __a=None	,    __a=None	,    __a=None	,    __a=None	,    )     ->		Any:







  """simple docstring"""





  if attention_mask is None:
    lowerCamelCase__:      Optional[Any]			  =np.where(input_ids != config.pad_token_id	,    1	,    0		)
  if decoder_attention_mask is None:
    lowerCamelCase__:      Dict			  =np.where(decoder_input_ids != config.pad_token_id	,    1	,    0		)
  if head_mask is None:
    lowerCamelCase__:      Optional[Any]			  =np.ones((config.encoder_layers, config.encoder_attention_heads)		)
  if decoder_head_mask is None:
    lowerCamelCase__:      Any			  =np.ones((config.decoder_layers, config.decoder_attention_heads)		)
  if cross_attn_head_mask is None:
    lowerCamelCase__:      List[str]			  =np.ones((config.decoder_layers, config.decoder_attention_heads)		)
  return {
      "input_ids": input_ids,
      "decoder_input_ids": decoder_input_ids,
      "attention_mask": attention_mask,
      "decoder_attention_mask": attention_mask,
  }



class    _SCREAMING_SNAKE_CASE						:







 '''simple docstring'''

 def __init__(self				:   Tuple       , UpperCAmelCase_				:   List[Any]       , UpperCAmelCase_				:   Dict=13       , UpperCAmelCase_				:   List[Any]=7       , UpperCAmelCase_				:   str=True       , UpperCAmelCase_				:   Optional[int]=False       , UpperCAmelCase_				:   Union[str, Any]=99       , UpperCAmelCase_				:   Any=16       , UpperCAmelCase_				:   Dict=2       , UpperCAmelCase_				:   Any=4       , UpperCAmelCase_				:   List[Any]=4       , UpperCAmelCase_				:   int="gelu"       , UpperCAmelCase_				:   int=0.1       , UpperCAmelCase_				:   str=0.1       , UpperCAmelCase_				:   Tuple=32       , UpperCAmelCase_				:   int=2       , UpperCAmelCase_				:   int=1       , UpperCAmelCase_				:   Union[str, Any]=0       , UpperCAmelCase_				:   Any=0.02       , )       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      int			  =parent
   lowerCamelCase__:      List[str]			  =batch_size
   lowerCamelCase__:      Optional[int]			  =seq_length
   lowerCamelCase__:      Optional[Any]			  =is_training
   lowerCamelCase__:      str			  =use_labels
   lowerCamelCase__:      Optional[Any]			  =vocab_size
   lowerCamelCase__:      int			  =hidden_size
   lowerCamelCase__:      Dict			  =num_hidden_layers
   lowerCamelCase__:      Any			  =num_attention_heads
   lowerCamelCase__:      str			  =intermediate_size
   lowerCamelCase__:      int			  =hidden_act
   lowerCamelCase__:      Tuple			  =hidden_dropout_prob
   lowerCamelCase__:      List[str]			  =attention_probs_dropout_prob
   lowerCamelCase__:      Optional[int]			  =max_position_embeddings
   lowerCamelCase__:      int			  =eos_token_id
   lowerCamelCase__:      Union[str, Any]			  =pad_token_id
   lowerCamelCase__:      List[str]			  =bos_token_id
   lowerCamelCase__:      int			  =initializer_range

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =np.clip(ids_tensor([self.batch_size, self.seq_length - 1]       , self.vocab_size)       , 3       , self.vocab_size)
   lowerCamelCase__:      str			  =np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)       , dtype=np.intaa))       , -1)

   lowerCamelCase__:      int			  =shift_tokens_right(UpperCAmelCase_       , 1       , 2)

   lowerCamelCase__:      Dict			  =BlenderbotConfig(
       vocab_size=self.vocab_size       , d_model=self.hidden_size       , encoder_layers=self.num_hidden_layers       , decoder_layers=self.num_hidden_layers       , encoder_attention_heads=self.num_attention_heads       , decoder_attention_heads=self.num_attention_heads       , encoder_ffn_dim=self.intermediate_size       , decoder_ffn_dim=self.intermediate_size       , dropout=self.hidden_dropout_prob       , attention_dropout=self.attention_probs_dropout_prob       , max_position_embeddings=self.max_position_embeddings       , eos_token_id=self.eos_token_id       , bos_token_id=self.bos_token_id       , pad_token_id=self.pad_token_id       , initializer_range=self.initializer_range       , use_cache=UpperCAmelCase_       , )
   lowerCamelCase__:      Any			  =prepare_blenderbot_inputs_dict(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_)
   return config, inputs_dict

 def 							SCREAMING_SNAKE_CASE_					(self				:   int)       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__  ,					lowerCamelCase__:      Dict			  =self.prepare_config_and_inputs()
   return config, inputs_dict

 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict       , UpperCAmelCase_				:   Optional[int]       , UpperCAmelCase_				:   List[str]       , UpperCAmelCase_				:   Dict)       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__:      Optional[Any]			  =20
   lowerCamelCase__:      Optional[int]			  =model_class_name(UpperCAmelCase_)

   lowerCamelCase__:      str			  =model.encode(inputs_dict["input_ids"])

   lowerCamelCase__  ,					lowerCamelCase__:      List[Any]			  =(
       inputs_dict["decoder_input_ids"],
       inputs_dict["decoder_attention_mask"],
   )

   lowerCamelCase__:      Union[str, Any]			  =model.init_cache(decoder_input_ids.shape[0]       , UpperCAmelCase_       , UpperCAmelCase_)
   lowerCamelCase__:      Union[str, Any]			  =jnp.ones((decoder_input_ids.shape[0], max_decoder_length)       , dtype="i4")

   lowerCamelCase__:      Tuple			  =jnp.broadcast_to(
       jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]       , (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)       , )
   lowerCamelCase__:      Union[str, Any]			  =model.decode(
       decoder_input_ids[:, :-1]       , UpperCAmelCase_       , decoder_attention_mask=UpperCAmelCase_       , past_key_values=UpperCAmelCase_       , decoder_position_ids=UpperCAmelCase_       , )

   lowerCamelCase__:      Union[str, Any]			  =jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]       , dtype="i4")
   lowerCamelCase__:      Dict			  =model.decode(
       decoder_input_ids[:, -1:]       , UpperCAmelCase_       , decoder_attention_mask=UpperCAmelCase_       , past_key_values=outputs_cache.past_key_values       , decoder_position_ids=UpperCAmelCase_       , )

   lowerCamelCase__:      List[Any]			  =model.decode(UpperCAmelCase_       , UpperCAmelCase_)

   lowerCamelCase__:      Optional[Any]			  =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
   self.parent.assertTrue(diff < 1E-3       , msg=F"""Max diff is {diff}""")







 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str]       , UpperCAmelCase_				:   Tuple       , UpperCAmelCase_				:   Optional[Any]       , UpperCAmelCase_				:   Tuple)       ->Optional[int]:



   '''simple docstring'''



   lowerCamelCase__:      List[str]			  =20
   lowerCamelCase__:      Optional[Any]			  =model_class_name(UpperCAmelCase_)

   lowerCamelCase__:      Any			  =model.encode(inputs_dict["input_ids"])

   lowerCamelCase__  ,					lowerCamelCase__:      Union[str, Any]			  =(
       inputs_dict["decoder_input_ids"],
       inputs_dict["decoder_attention_mask"],
   )

   lowerCamelCase__:      Optional[int]			  =jnp.concatenate(
       [
           decoder_attention_mask,
           jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])),
       ]       , axis=-1       , )

   lowerCamelCase__:      Union[str, Any]			  =model.init_cache(decoder_input_ids.shape[0]       , UpperCAmelCase_       , UpperCAmelCase_)
   lowerCamelCase__:      Tuple			  =jnp.broadcast_to(
       jnp.arange(decoder_input_ids.shape[-1] - 1)[None, :]       , (decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)       , )

   lowerCamelCase__:      List[Any]			  =model.decode(
       decoder_input_ids[:, :-1]       , UpperCAmelCase_       , decoder_attention_mask=UpperCAmelCase_       , past_key_values=UpperCAmelCase_       , decoder_position_ids=UpperCAmelCase_       , )
   lowerCamelCase__:      Dict			  =jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]       , dtype="i4")
   lowerCamelCase__:      str			  =model.decode(
       decoder_input_ids[:, -1:]       , UpperCAmelCase_       , past_key_values=outputs_cache.past_key_values       , decoder_attention_mask=UpperCAmelCase_       , decoder_position_ids=UpperCAmelCase_       , )

   lowerCamelCase__:      Union[str, Any]			  =model.decode(UpperCAmelCase_       , UpperCAmelCase_       , decoder_attention_mask=UpperCAmelCase_)

   lowerCamelCase__:      str			  =np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])))
   self.parent.assertTrue(diff < 1E-3       , msg=F"""Max diff is {diff}""")



@require_flax
class    _SCREAMING_SNAKE_CASE						(	unittest.TestCase    ):







 '''simple docstring'''
 lowercase_		=			99

 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->int:



   '''simple docstring'''



   lowerCamelCase__:      Union[str, Any]			  =np.array(
       [
           [71, 82, 18, 33, 46, 91, 2],
           [68, 34, 26, 58, 30, 82, 2],
           [5, 97, 17, 39, 94, 40, 2],
           [76, 83, 94, 25, 70, 78, 2],
           [87, 59, 41, 35, 48, 66, 2],
           [55, 13, 16, 58, 5, 2, 1],  # note padding
           [64, 27, 31, 51, 12, 75, 2],
           [52, 64, 86, 17, 83, 39, 2],
           [48, 61, 9, 24, 71, 82, 2],
           [26, 1, 60, 48, 22, 13, 2],
           [21, 5, 62, 28, 14, 76, 2],
           [45, 98, 37, 86, 59, 48, 2],
           [70, 70, 50, 9, 28, 0, 2],
       ]       , dtype=np.intaa       , )

   lowerCamelCase__:      Optional[Any]			  =input_ids.shape[0]
   lowerCamelCase__:      List[str]			  =BlenderbotConfig(
       vocab_size=self.vocab_size       , d_model=24       , encoder_layers=2       , decoder_layers=2       , encoder_attention_heads=2       , decoder_attention_heads=2       , encoder_ffn_dim=32       , decoder_ffn_dim=32       , max_position_embeddings=48       , eos_token_id=2       , pad_token_id=1       , bos_token_id=0       , )
   return config, input_ids, batch_size

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->Optional[Any]:



   '''simple docstring'''



   lowerCamelCase__  ,					lowerCamelCase__  ,					lowerCamelCase__:      Any			  =self._get_config_and_data()
   lowerCamelCase__:      Dict			  =FlaxBlenderbotForConditionalGeneration(UpperCAmelCase_)
   lowerCamelCase__:      Dict			  =lm_model(input_ids=UpperCAmelCase_)
   lowerCamelCase__:      Dict			  =(batch_size, input_ids.shape[1], config.vocab_size)
   self.assertEqual(outputs["logits"].shape       , UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Dict)       ->str:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =BlenderbotConfig(
       vocab_size=self.vocab_size       , d_model=14       , encoder_layers=2       , decoder_layers=2       , encoder_attention_heads=2       , decoder_attention_heads=2       , encoder_ffn_dim=8       , decoder_ffn_dim=8       , max_position_embeddings=48       , )
   lowerCamelCase__:      str			  =FlaxBlenderbotForConditionalGeneration(UpperCAmelCase_)
   lowerCamelCase__:      Optional[int]			  =np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]       , dtype=np.intaa)
   lowerCamelCase__:      Optional[int]			  =np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]       , dtype=np.intaa)
   lowerCamelCase__:      List[str]			  =lm_model(input_ids=UpperCAmelCase_       , decoder_input_ids=UpperCAmelCase_)
   lowerCamelCase__:      Optional[int]			  =(*summary.shape, config.vocab_size)
   self.assertEqual(outputs["logits"].shape       , UpperCAmelCase_)







 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Tuple:



   '''simple docstring'''



   lowerCamelCase__:      Optional[int]			  =np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]       , dtype=np.intaa)
   lowerCamelCase__:      Optional[int]			  =shift_tokens_right(UpperCAmelCase_       , 1       , 2)
   lowerCamelCase__:      List[str]			  =np.equal(UpperCAmelCase_       , 1).astype(np.floataa).sum()
   lowerCamelCase__:      Tuple			  =np.equal(UpperCAmelCase_       , 1).astype(np.floataa).sum()
   self.assertEqual(shifted.shape       , input_ids.shape)
   self.assertEqual(UpperCAmelCase_       , n_pad_before - 1)
   self.assertTrue(np.equal(shifted[:, 0]       , 2).all())



@require_flax
class    _SCREAMING_SNAKE_CASE						(	__SCREAMING_SNAKE_CASE    ,							unittest.TestCase    ,							__SCREAMING_SNAKE_CASE    ):







 '''simple docstring'''
 lowercase_		=			True
 lowercase_		=			(
     (
         FlaxBlenderbotModel,
         FlaxBlenderbotForConditionalGeneration,
     )
     if is_flax_available()
     else ()
 )
 lowercase_		=			(FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[str])       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__:      List[Any]			  =FlaxBlenderbotModelTester(self)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->List[str]:



   '''simple docstring'''



   lowerCamelCase__  ,					lowerCamelCase__:      List[str]			  =self.model_tester.prepare_config_and_inputs()
   for model_class in self.all_model_classes:
     self.model_tester.check_use_cache_forward(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple)       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__  ,					lowerCamelCase__:      List[str]			  =self.model_tester.prepare_config_and_inputs()
   for model_class in self.all_model_classes:
     self.model_tester.check_use_cache_forward_with_attn_mask(UpperCAmelCase_       , UpperCAmelCase_       , UpperCAmelCase_)

 def 							SCREAMING_SNAKE_CASE_					(self				:   Union[str, Any])       ->str:



   '''simple docstring'''



   lowerCamelCase__  ,					lowerCamelCase__:      Union[str, Any]			  =self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__):
       lowerCamelCase__:      List[str]			  =self._prepare_for_class(UpperCAmelCase_       , UpperCAmelCase_)
       lowerCamelCase__:      Optional[int]			  =model_class(UpperCAmelCase_)

       @jax.jit
       def encode_jitted(UpperCAmelCase_				:   List[Any]       , UpperCAmelCase_				:   Any=None       , **UpperCAmelCase_				:   List[str]):
         return model.encode(input_ids=UpperCAmelCase_       , attention_mask=UpperCAmelCase_)

       with self.subTest("JIT Enabled"):
         lowerCamelCase__:      Any			  =encode_jitted(**UpperCAmelCase_).to_tuple()

       with self.subTest("JIT Disabled"):
         with jax.disable_jit():
           lowerCamelCase__:      Tuple			  =encode_jitted(**UpperCAmelCase_).to_tuple()

       self.assertEqual(len(UpperCAmelCase_)       , len(UpperCAmelCase_))
       for jitted_output, output in zip(UpperCAmelCase_       , UpperCAmelCase_):
         self.assertEqual(jitted_output.shape       , output.shape)

 def 							SCREAMING_SNAKE_CASE_					(self				:   List[Any])       ->List[Any]:



   '''simple docstring'''



   lowerCamelCase__  ,					lowerCamelCase__:      List[Any]			  =self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
     with self.subTest(model_class.__name__):
       lowerCamelCase__:      Optional[Any]			  =model_class(UpperCAmelCase_)
       lowerCamelCase__:      List[Any]			  =model.encode(inputs_dict["input_ids"]       , inputs_dict["attention_mask"])

       lowerCamelCase__:      int			  ={
           "decoder_input_ids": inputs_dict["decoder_input_ids"],
           "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
           "encoder_outputs": encoder_outputs,
       }

       @jax.jit
       def decode_jitted(UpperCAmelCase_				:   Optional[Any]       , UpperCAmelCase_				:   Dict       , UpperCAmelCase_				:   Optional[int]):
         return model.decode(
             decoder_input_ids=UpperCAmelCase_       , decoder_attention_mask=UpperCAmelCase_       , encoder_outputs=UpperCAmelCase_       , )

       with self.subTest("JIT Enabled"):
         lowerCamelCase__:      int			  =decode_jitted(**UpperCAmelCase_).to_tuple()

       with self.subTest("JIT Disabled"):
         with jax.disable_jit():
           lowerCamelCase__:      int			  =decode_jitted(**UpperCAmelCase_).to_tuple()

       self.assertEqual(len(UpperCAmelCase_)       , len(UpperCAmelCase_))
       for jitted_output, output in zip(UpperCAmelCase_       , UpperCAmelCase_):
         self.assertEqual(jitted_output.shape       , output.shape)

 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   Any)       ->Union[str, Any]:



   '''simple docstring'''



   for model_class_name in self.all_model_classes:
     lowerCamelCase__:      Optional[int]			  =model_class_name.from_pretrained("facebook/blenderbot-400M-distill")
     # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
     lowerCamelCase__:      int			  =np.ones((1, 1)) * model.config.eos_token_id
     lowerCamelCase__:      str			  =model(UpperCAmelCase_)
     self.assertIsNotNone(UpperCAmelCase_)







 @unittest.skipUnless(jax_device != "cpu"       , "3B test too slow on CPU.")
 @slow
 def 							SCREAMING_SNAKE_CASE_					(self				:   Tuple)       ->Dict:



   '''simple docstring'''



   lowerCamelCase__:      Dict			  ={"num_beams": 1, "early_stopping": True, "min_length": 15, "max_length": 25}
   lowerCamelCase__:      Union[str, Any]			  ={"skip_special_tokens": True, "clean_up_tokenization_spaces": True}

   lowerCamelCase__:      Dict			  =FlaxBlenderbotForConditionalGeneration.from_pretrained("facebook/blenderbot-3B"       , from_pt=UpperCAmelCase_)
   lowerCamelCase__:      List[str]			  =BlenderbotTokenizer.from_pretrained("facebook/blenderbot-3B")

   lowerCamelCase__:      Any			  =["Sam"]
   lowerCamelCase__:      Tuple			  =tokenizer(UpperCAmelCase_       , return_tensors="jax")

   lowerCamelCase__:      Optional[Any]			  =model.generate(**UpperCAmelCase_       , **UpperCAmelCase_)
   lowerCamelCase__:      Any			  ="Sam is a great name. It means \"sun\" in Gaelic."

   lowerCamelCase__:      Optional[Any]			  =tokenizer.batch_decode(UpperCAmelCase_       , **UpperCAmelCase_)
   assert generated_txt[0].strip() == tgt_text

