


from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__UpperCamelCase     		=							logging.get_logger(__name__)

__UpperCamelCase     		=							{
    'xlm-mlm-en-2048': 'https://huggingface.co/xlm-mlm-en-2048/resolve/main/config.json',
    'xlm-mlm-ende-1024': 'https://huggingface.co/xlm-mlm-ende-1024/resolve/main/config.json',
    'xlm-mlm-enfr-1024': 'https://huggingface.co/xlm-mlm-enfr-1024/resolve/main/config.json',
    'xlm-mlm-enro-1024': 'https://huggingface.co/xlm-mlm-enro-1024/resolve/main/config.json',
    'xlm-mlm-tlm-xnli15-1024': 'https://huggingface.co/xlm-mlm-tlm-xnli15-1024/resolve/main/config.json',
    'xlm-mlm-xnli15-1024': 'https://huggingface.co/xlm-mlm-xnli15-1024/resolve/main/config.json',
    'xlm-clm-enfr-1024': 'https://huggingface.co/xlm-clm-enfr-1024/resolve/main/config.json',
    'xlm-clm-ende-1024': 'https://huggingface.co/xlm-clm-ende-1024/resolve/main/config.json',
    'xlm-mlm-17-1280': 'https://huggingface.co/xlm-mlm-17-1280/resolve/main/config.json',
    'xlm-mlm-100-1280': 'https://huggingface.co/xlm-mlm-100-1280/resolve/main/config.json',
}







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			_UpperCamelCase		:       Optional[int] 			=		'xlm'
			_UpperCamelCase		:       str 			=		{
			    'hidden_size': 'emb_dim',
			    'num_attention_heads': 'n_heads',
			    'num_hidden_layers': 'n_layers',
			    'n_words': 'vocab_size',  # For backward compatibility
			}




			def __init__(				self   ,    snake_case=30145   ,    snake_case=2048   ,    snake_case=12   ,    snake_case=16   ,    snake_case=0.1   ,    snake_case=0.1   ,    snake_case=True   ,    snake_case=False   ,    snake_case=False   ,    snake_case=False   ,    snake_case=1   ,    snake_case=True   ,    snake_case=512   ,    snake_case=2048**-0.5   ,    snake_case=1E-1_2   ,    snake_case=0.02   ,    snake_case=0   ,    snake_case=1   ,    snake_case=2   ,    snake_case=3   ,    snake_case=5   ,    snake_case=True   ,    snake_case="first"   ,    snake_case=True   ,    snake_case=None   ,    snake_case=True   ,    snake_case=0.1   ,    snake_case=5   ,    snake_case=5   ,    snake_case=0   ,    snake_case=0   ,    snake_case=2   ,    snake_case=0   ,    **snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							vocab_size
						UpperCamelCase__       =							emb_dim
						UpperCamelCase__       =							n_layers
						UpperCamelCase__       =							n_heads
						UpperCamelCase__       =							dropout
						UpperCamelCase__       =							attention_dropout
						UpperCamelCase__       =							gelu_activation
						UpperCamelCase__       =							sinusoidal_embeddings
						UpperCamelCase__       =							causal
						UpperCamelCase__       =							asm
						UpperCamelCase__       =							n_langs
						UpperCamelCase__       =							use_lang_emb
						UpperCamelCase__       =							layer_norm_eps
						UpperCamelCase__       =							bos_index
						UpperCamelCase__       =							eos_index
						UpperCamelCase__       =							pad_index
						UpperCamelCase__       =							unk_index
						UpperCamelCase__       =							mask_index
						UpperCamelCase__       =							is_encoder
						UpperCamelCase__       =							max_position_embeddings
						UpperCamelCase__       =							embed_init_std
						UpperCamelCase__       =							init_std
						UpperCamelCase__       =							summary_type
						UpperCamelCase__       =							summary_use_proj
						UpperCamelCase__       =							summary_activation
						UpperCamelCase__       =							summary_proj_to_labels
						UpperCamelCase__       =							summary_first_dropout
						UpperCamelCase__       =							start_n_top
						UpperCamelCase__       =							end_n_top
						UpperCamelCase__       =							mask_token_id
						UpperCamelCase__       =							lang_id

						if "n_words" in kwargs:
									UpperCamelCase__       =							kwargs["n_words"]

						super().__init__(pad_token_id=snake_case   ,    bos_token_id=snake_case   ,    **snake_case  )







class 							lowerCamelCase__						(				UpperCAmelCase   ):






			"""simple docstring"""






			@property
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						if self.task == "multiple-choice":
									UpperCamelCase__       =							{0: "batch", 1: "choice", 2: "sequence"}
						else:
									UpperCamelCase__       =							{0: "batch", 1: "sequence"}
						return OrderedDict(
						    [
						        ("input_ids", dynamic_axis),
						        ("attention_mask", dynamic_axis),
						        ("token_type_ids", dynamic_axis),
						    ]  )






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  )


