



'''simple docstring'''





import unittest
from pathlib import Path
from tempfile import NamedTemporaryFile, TemporaryDirectory

from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
from transformers.convert_graph_to_onnx import (
    convert,
    ensure_valid_input,
    generate_identified_filename,
    infer_shapes,
    quantize,
)
from transformers.testing_utils import require_tf, require_tokenizers, require_torch, slow


class 					snake_case			:



       def 				_a   (					self  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_   )	->				Union[str, Any]:
           return None





class 					snake_case			:



       def 				_a   (					self  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_   )	->				Union[str, Any]:
           return None





class 					snake_case			(unittest.TestCase ):
       lowerCAmelCase__	:str     =       [
           # (model_name, model_kwargs)
           ('bert-base-cased', {}),
           ('gpt2', {'use_cache': False}),  # We don't support exporting GPT2 past keys anymore
       ]



       @require_tf
       @slow
       def 				_a   (					self   )	->				str:


           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
               self._test_export(UpperCAmelCase_  ,"tf"  ,12  ,**UpperCAmelCase_   )



       @require_torch
       @slow
       def 				_a   (					self   )	->				Dict:


           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
               self._test_export(UpperCAmelCase_  ,"pt"  ,12  ,**UpperCAmelCase_   )



       @require_torch
       @slow
       def 				_a   (					self   )	->				List[Any]:
           from transformers import BertModel

           lowercase__									=    ["[UNK]", "[SEP]", "[CLS]", "[PAD]", "[MASK]", "some", "other", "words"]
           with NamedTemporaryFile(mode="w+t"   ) as vocab_file:
               vocab_file.write("\n".join(UpperCAmelCase_   )   )
               vocab_file.flush()
               lowercase__									=    BertTokenizerFast(vocab_file.name   )

           with TemporaryDirectory() as bert_save_dir:
               lowercase__									=    BertModel(BertConfig(vocab_size=len(UpperCAmelCase_   )   )   )
               model.save_pretrained(UpperCAmelCase_   )
               self._test_export(UpperCAmelCase_  ,"pt"  ,12  ,UpperCAmelCase_   )



       @require_tf
       @slow
       def 				_a   (					self   )	->				List[Any]:


           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
               lowercase__									=    self._test_export(UpperCAmelCase_  ,"tf"  ,12  ,**UpperCAmelCase_   )
               lowercase__									=    quantize(Path(UpperCAmelCase_   )   )

               # Ensure the actual quantized model is not bigger than the original one
               if quantized_path.stat().st_size >= Path(UpperCAmelCase_   ).stat().st_size:
                   self.fail("Quantized model is bigger than initial ONNX model"   )



       @require_torch
       @slow
       def 				_a   (					self   )	->				Tuple:


           for model, model_kwargs in OnnxExportTestCase.MODEL_TO_TEST:
               lowercase__									=    self._test_export(UpperCAmelCase_  ,"pt"  ,12  ,**UpperCAmelCase_   )
               lowercase__									=    quantize(UpperCAmelCase_   )

               # Ensure the actual quantized model is not bigger than the original one
               if quantized_path.stat().st_size >= Path(UpperCAmelCase_   ).stat().st_size:
                   self.fail("Quantized model is bigger than initial ONNX model"   )



       def 				_a   (					self  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_=None  ,**UpperCAmelCase_   )	->				Optional[int]:


           try:
               # Compute path
               with TemporaryDirectory() as tempdir:
                   lowercase__									=    Path(UpperCAmelCase_   ).joinpath("model.onnx"   )

               # Remove folder if exists
               if path.parent.exists():
                   path.parent.rmdir()

               # Export
               convert(UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_  ,**UpperCAmelCase_   )

               return path
           except Exception as e:
               self.fail(UpperCAmelCase_   )



       @require_torch
       @require_tokenizers
       @slow
       def 				_a   (					self   )	->				str:
           from transformers import BertModel

           lowercase__									=    BertModel(BertConfig.from_pretrained("lysandre/tiny-bert-random"   )   )
           lowercase__									=    BertTokenizerFast.from_pretrained("lysandre/tiny-bert-random"   )
           self._test_infer_dynamic_axis(UpperCAmelCase_  ,UpperCAmelCase_  ,"pt"   )



       @require_tf
       @require_tokenizers
       @slow
       def 				_a   (					self   )	->				Union[str, Any]:
           from transformers import TFBertModel

           lowercase__									=    TFBertModel(BertConfig.from_pretrained("lysandre/tiny-bert-random"   )   )
           lowercase__									=    BertTokenizerFast.from_pretrained("lysandre/tiny-bert-random"   )
           self._test_infer_dynamic_axis(UpperCAmelCase_  ,UpperCAmelCase_  ,"tf"   )



       def 				_a   (					self  ,UpperCAmelCase_  ,UpperCAmelCase_  ,UpperCAmelCase_   )	->				Any:
           lowercase__									=    FeatureExtractionPipeline(UpperCAmelCase_  ,UpperCAmelCase_   )

           lowercase__									=    ["input_ids", "token_type_ids", "attention_mask", "output_0", "output_1"]
           lowercase__							,							lowercase__							,							lowercase__							,							lowercase__									=    infer_shapes(UpperCAmelCase_  ,UpperCAmelCase_   )

           # Assert all variables are present
           self.assertEqual(len(UpperCAmelCase_   )  ,len(UpperCAmelCase_   )   )
           self.assertTrue(all(var_name in shapes for var_name in variable_names   )   )
           self.assertSequenceEqual(variable_names[:3]  ,UpperCAmelCase_   )
           self.assertSequenceEqual(variable_names[3:]  ,UpperCAmelCase_   )

           # Assert inputs are {0: batch, 1: sequence}
           for var_name in ["input_ids", "token_type_ids", "attention_mask"]:
               self.assertDictEqual(shapes[var_name]  ,{0: "batch", 1: "sequence"}   )

           # Assert outputs are {0: batch, 1: sequence} and {0: batch}
           self.assertDictEqual(shapes["output_0"]  ,{0: "batch", 1: "sequence"}   )
           self.assertDictEqual(shapes["output_1"]  ,{0: "batch"}   )



       def 				_a   (					self   )	->				Union[str, Any]:
           lowercase__									=    ["input_ids", "attention_mask", "token_type_ids"]
           lowercase__									=    {"input_ids": [1, 2, 3, 4], "attention_mask": [0, 0, 0, 0], "token_type_ids": [1, 1, 1, 1]}
           lowercase__							,							lowercase__									=    ensure_valid_input(FuncContiguousArgs()  ,UpperCAmelCase_  ,UpperCAmelCase_   )

           # Should have exactly the same number of args (all are valid)
           self.assertEqual(len(UpperCAmelCase_   )  ,3   )

           # Should have exactly the same input names
           self.assertEqual(set(UpperCAmelCase_   )  ,set(UpperCAmelCase_   )   )

           # Parameter should be reordered according to their respective place in the function:
           # (input_ids, token_type_ids, attention_mask)
           self.assertEqual(UpperCAmelCase_  ,(tokens["input_ids"], tokens["token_type_ids"], tokens["attention_mask"])   )

           # Generated args are interleaved with another args (for instance parameter "past" in GPT2)
           lowercase__							,							lowercase__									=    ensure_valid_input(FuncNonContiguousArgs()  ,UpperCAmelCase_  ,UpperCAmelCase_   )

           # Should have exactly the one arg (all before the one not provided "some_other_args")
           self.assertEqual(len(UpperCAmelCase_   )  ,1   )
           self.assertEqual(len(UpperCAmelCase_   )  ,1   )

           # Should have only "input_ids"
           self.assertEqual(inputs_args[0]  ,tokens["input_ids"]   )
           self.assertEqual(ordered_input_names[0]  ,"input_ids"   )







       def 				_a   (					self   )	->				List[str]:
           lowercase__									=    generate_identified_filename(Path("/home/something/my_fake_model.onnx"   )  ,"-test"   )
           self.assertEqual("/home/something/my_fake_model-test.onnx"  ,generated.as_posix()   )








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)
__UpperCamelCase     		=							logging.getLogger()

__UpperCamelCase     		=							logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			def 					snake_case__					(				self   ,    snake_case  ):

						'''simple docstring'''


						os.makedirs(snake_case   ,    exist_ok=snake_case  )
						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(snake_case   ,    F'''{split}.{field}'''  )   ,    "w"  ) as f:
															f.write(snake_case  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case = "pytorch"  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.get_auto_remove_tmp_dir()
						UpperCamelCase__       =							os.path.join(snake_case   ,    "output"  )
						UpperCamelCase__       =							os.path.join(snake_case   ,    "data"  )
						self._create_dummy_data(data_dir=snake_case  )

						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(snake_case   ,    env=self.get_env()  )

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




			@require_torch_gpu
			def 					snake_case__					(				self  ):

						'''simple docstring'''


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




			@require_torch_multi_gpu
			def 					snake_case__					(				self  ):

						'''simple docstring'''


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




			@require_torch_gpu
			@require_ray
			def 					snake_case__					(				self  ):

						'''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 					snake_case__					(				self  ):

						'''simple docstring'''


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


