






import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
__a:     Dict   					=					logging.getLogger()

__a:     Optional[int]   					=					logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
class 		SCREAMING_SNAKE_CASE__	(  __lowerCAmelCase     ):

  '''simple docstring'''







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


         """simple docstring"""
         os.makedirs(lowerCamelCase_       ,				exist_ok=lowerCamelCase_							)
         _UpperCAmelCase							=    {"""source""": """What is love ?""", """target""": """life"""}
         _UpperCAmelCase							=    {"""train""": 12, """val""": 2, """test""": 2}
         for split in ["train", "test", "val"]:
                for field in ["source", "target"]:
                       _UpperCAmelCase							=    """\n""".join([contents[field]] * n_lines[split]							)
                       with open(os.path.join(lowerCamelCase_       ,				f"""{split}.{field}"""							)       ,				"""w"""							) as f:
                              f.write(lowerCamelCase_							)


  def  lowerCamelCase       (      self		:			Tuple       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			str = "pytorch"							)						->	Dict:


         """simple docstring"""
         _UpperCAmelCase							=    self.get_auto_remove_tmp_dir()
         _UpperCAmelCase							=    os.path.join(lowerCamelCase_       ,				"""output"""							)
         _UpperCAmelCase							=    os.path.join(lowerCamelCase_       ,				"""data"""							)
         self._create_dummy_data(data_dir=lowerCamelCase_							)

         _UpperCAmelCase							=    f"""
                --data_dir {data_dir} \
                --output_dir {output_dir} \
                --model_name_or_path facebook/rag-sequence-base \
                --model_type rag_sequence \
                --do_train \
                --do_predict \
                --n_val -1 \
                --val_check_interval 1.0 \
                --train_batch_size 2 \
                --eval_batch_size 1 \
                --max_source_length 25 \
                --max_target_length 25 \
                --val_max_target_length 25 \
                --test_max_target_length 25 \
                --label_smoothing 0.1 \
                --dropout 0.1 \
                --attention_dropout 0.1 \
                --weight_decay 0.001 \
                --adam_epsilon 1e-08 \
                --max_grad_norm 0.1 \
                --lr_scheduler polynomial \
                --learning_rate 3e-04 \
                --num_train_epochs 1 \
                --warmup_steps 4 \
                --gradient_accumulation_steps 1 \
                --distributed-port 8787 \
                --use_dummy_dataset 1 \
                --distributed_retriever {distributed_retriever} \
            """.split()

         if gpus > 0:
                testargs.append(f"""--gpus={gpus}"""							)
                if is_apex_available():
                       testargs.append("""--fp16"""							)
         else:
                testargs.append("""--gpus=0"""							)
                testargs.append("""--distributed_backend=ddp_cpu"""							)
                testargs.append("""--num_processes=2"""							)

         _UpperCAmelCase							=    [sys.executable, str(Path(finetune_rag.__file__							).resolve()							)] + testargs
         execute_subprocess_async(lowerCamelCase_       ,				env=self.get_env()							)

         _UpperCAmelCase							=    os.path.join(lowerCamelCase_       ,				"""metrics.json"""							)
         with open(lowerCamelCase_							) as f:
                _UpperCAmelCase							=    json.load(lowerCamelCase_							)
         return result


  @require_torch_gpu
  def  lowerCamelCase       (      self		:			Tuple							)						->	Union[str, Any]:


         """simple docstring"""
         _UpperCAmelCase							=    self._run_finetune(gpus=1							)
         self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]       ,				0.2							)


  @require_torch_multi_gpu
  def  lowerCamelCase       (      self		:			List[Any]							)						->	str:


         """simple docstring"""
         _UpperCAmelCase							=    self._run_finetune(gpus=2							)
         self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]       ,				0.2							)


  @require_torch_gpu
  @require_ray
  def  lowerCamelCase       (      self		:			Union[str, Any]							)						->	Optional[int]:


         """simple docstring"""
         _UpperCAmelCase							=    self._run_finetune(gpus=1       ,				distributed_retriever="""ray"""							)
         self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]       ,				0.2							)





  @require_torch_multi_gpu
  @require_ray
  def  lowerCamelCase       (      self		:			int							)						->	Any:


         """simple docstring"""
         _UpperCAmelCase							=    self._run_finetune(gpus=1       ,				distributed_retriever="""ray"""							)
         self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]       ,				0.2							)
import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
     import os

     import jax.numpy as jnp
     from jax import jit

     from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
     from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

     _SCREAMING_SNAKE_CASE 		=      """0.12"""  # assumed parallelism: 8


if is_torch_available():
     import torch





def 			lowercase(	UpperCamelCase_							,					UpperCamelCase_							,					UpperCamelCase_=None       )		->	int:







 '''simple docstring'''
 if rng is None:
  UpperCamelCase			  =      random.Random()

 UpperCamelCase			  =      1
 for dim in shape:
  total_dims *= dim

 UpperCamelCase			  =      []
 for _ in range(UpperCamelCase_       ):
  values.append(rng.randint(0							,					vocab_size - 1       )       )

 UpperCamelCase			  =      np.array(UpperCamelCase_							,					dtype=jnp.intaa       ).reshape(UpperCamelCase_       )

 return output





def 			lowercase(	UpperCamelCase_							,					UpperCamelCase_=None       )		->	List[str]:







 '''simple docstring'''
 UpperCamelCase			  =      ids_tensor(UpperCamelCase_							,					vocab_size=2							,					rng=UpperCamelCase_       )
 # make sure that at least one token is attended to for each batch
 UpperCamelCase			  =      1
 return attn_mask
@require_flax
class SCREAMING_SNAKE_CASE_		:
   __lowerCAmelCase           =				None
   __lowerCAmelCase           =				()







   def 						lowerCamelCase_	(     self						:						Tuple				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase			  =      self.model_tester.prepare_config_and_inputs_for_common()

    # cut to half length & take max batch_size 3
    UpperCamelCase			  =      2
    UpperCamelCase			  =      inputs["""input_ids"""].shape[-1] // 2
    UpperCamelCase			  =      inputs["""input_ids"""][:max_batch_size, :sequence_length]

    UpperCamelCase			  =      jnp.ones_like(lowerCamelCase_				)
    UpperCamelCase			  =      attention_mask[:max_batch_size, :sequence_length]

    # generate max 5 tokens
    UpperCamelCase			  =      input_ids.shape[-1] + 5
    if config.eos_token_id is not None and config.pad_token_id is None:
     # hack to allow generate for models such as GPT2 as is done in `generate()`
     UpperCamelCase			  =      config.eos_token_id
    return config, input_ids, attention_mask, max_length







   @is_pt_flax_cross_test
   def 						lowerCamelCase_	(     self						:						Dict				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      False
    UpperCamelCase			  =      max_length
    UpperCamelCase			  =      0

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model_class.__name__[4:]  # Skip the "Flax" at the beginning
     UpperCamelCase			  =      getattr(lowerCamelCase_   ,   lowerCamelCase_				)
     UpperCamelCase			  =      pt_model_class(lowerCamelCase_				).eval()
     UpperCamelCase			  =      load_flax_weights_in_pytorch_model(lowerCamelCase_   ,   flax_model.params				)

     UpperCamelCase			  =      flax_model.generate(lowerCamelCase_				).sequences
     UpperCamelCase			  =      pt_model.generate(torch.tensor(lowerCamelCase_   ,   dtype=torch.long				)				)

     if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
      UpperCamelCase			  =      flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

     self.assertListEqual(pt_generation_outputs.numpy().tolist()   ,   flax_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						List[str]				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      False
    UpperCamelCase			  =      max_length

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						List[Any]				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      True
    UpperCamelCase			  =      max_length

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						Any				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      False
    UpperCamelCase			  =      max_length
    UpperCamelCase			  =      2

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						int				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      False
    UpperCamelCase			  =      max_length
    UpperCamelCase			  =      2
    UpperCamelCase			  =      2

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[0]   ,   input_ids.shape[0] * config.num_return_sequences				)







   def 						lowerCamelCase_	(     self						:						Union[str, Any]				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      True
    UpperCamelCase			  =      max_length
    UpperCamelCase			  =      0.8
    UpperCamelCase			  =      10
    UpperCamelCase			  =      0.3
    UpperCamelCase			  =      1
    UpperCamelCase			  =      8
    UpperCamelCase			  =      9

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						Optional[int]				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      max_length
    UpperCamelCase			  =      1
    UpperCamelCase			  =      8
    UpperCamelCase			  =      9

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						Any				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()
    UpperCamelCase			  =      max_length
    UpperCamelCase			  =      2
    UpperCamelCase			  =      1
    UpperCamelCase			  =      8
    UpperCamelCase			  =      9

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						int				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()

    # pad attention mask on the left
    UpperCamelCase			  =      attention_mask.at[(0, 0)].set(0				)

    UpperCamelCase			  =      False
    UpperCamelCase			  =      max_length

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_   ,   attention_mask=lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_   ,   attention_mask=lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						Union[str, Any]				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()

    # pad attention mask on the left
    UpperCamelCase			  =      attention_mask.at[(0, 0)].set(0				)

    UpperCamelCase			  =      True
    UpperCamelCase			  =      max_length

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_   ,   attention_mask=lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_   ,   attention_mask=lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)







   def 						lowerCamelCase_	(     self						:						Any				):





    """simple docstring"""


    UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase  ,     UpperCamelCase			  =      self._get_input_ids_and_config()

    # pad attention mask on the left
    UpperCamelCase			  =      attention_mask.at[(0, 0)].set(0				)

    UpperCamelCase			  =      2
    UpperCamelCase			  =      max_length

    for model_class in self.all_generative_model_classes:
     UpperCamelCase			  =      model_class(lowerCamelCase_				)

     UpperCamelCase			  =      model.generate(lowerCamelCase_   ,   attention_mask=lowerCamelCase_				).sequences
     self.assertEqual(generation_outputs.shape[-1]   ,   lowerCamelCase_				)

     UpperCamelCase			  =      jit(model.generate				)
     UpperCamelCase			  =      jit_generate(lowerCamelCase_   ,   attention_mask=lowerCamelCase_				).sequences

     self.assertListEqual(generation_outputs.tolist()   ,   jit_generation_outputs.tolist()				)





@require_flax
class SCREAMING_SNAKE_CASE_		(      unittest.TestCase		):







   def 						lowerCamelCase_	(     self						:						Union[str, Any]				):





    """simple docstring"""


    UpperCamelCase			  =      AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-bert"""				)
    UpperCamelCase			  =      FlaxAutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-bert-flax-only"""				)

    UpperCamelCase			  =      """Hello world"""
    UpperCamelCase			  =      tokenizer(lowerCamelCase_   ,   return_tensors="""np"""				).input_ids

    # typos are quickly detected (the correct argument is `do_sample`)
    with self.assertRaisesRegex(lowerCamelCase_   ,   """do_samples"""				):
     model.generate(lowerCamelCase_   ,   do_samples=lowerCamelCase_				)

    # arbitrary arguments that will not be used anywhere are also not accepted
    with self.assertRaisesRegex(lowerCamelCase_   ,   """foo"""				):
     UpperCamelCase			  =      {"""foo""": """bar"""}
     model.generate(lowerCamelCase_   ,   **lowerCamelCase_				)
