





"""simple docstring"""




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 			_a						(				_lowerCAmelCase    ):






				def     snake_case     (			self		:   Optional[int],			lowerCAmelCase__		:   Dict							)     ->				Union[str, Any]:




										'''simple docstring'''



										os.makedirs(lowerCAmelCase__,			exist_ok=lowerCAmelCase__							)
										_UpperCamelCase      : Optional[Any]         =   {'''source''': '''What is love ?''', '''target''': '''life'''}
										_UpperCamelCase      : Optional[int]         =   {'''train''': 1_2, '''val''': 2, '''test''': 2}
										for split in ["train", "test", "val"]:
																for field in ["source", "target"]:
																						_UpperCamelCase      : int         =   '''\n'''.join([contents[field]] * n_lines[split]							)
																						with open(os.path.join(lowerCAmelCase__,			f"""{split}.{field}"""							),			'''w'''							) as f:
																												f.write(lowerCAmelCase__							)






				def     snake_case     (			self		:   Optional[Any],			lowerCAmelCase__		:   int,			lowerCAmelCase__		:   str = "pytorch"							)     ->				str:




										'''simple docstring'''



										_UpperCamelCase      : Any         =   self.get_auto_remove_tmp_dir()
										_UpperCamelCase      : List[Any]         =   os.path.join(lowerCAmelCase__,			'''output'''							)
										_UpperCamelCase      : Any         =   os.path.join(lowerCAmelCase__,			'''data'''							)
										self._create_dummy_data(data_dir=lowerCAmelCase__							)

										_UpperCamelCase      : Any         =   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      : Dict         =   [sys.executable, str(Path(finetune_rag.__file__							).resolve()							)] + testargs
										execute_subprocess_async(lowerCAmelCase__,			env=self.get_env()							)

										_UpperCamelCase      : Optional[int]         =   os.path.join(lowerCAmelCase__,			'''metrics.json'''							)
										with open(lowerCAmelCase__							) as f:
																_UpperCamelCase      : Union[str, Any]         =   json.load(lowerCAmelCase__							)
										return result






				@require_torch_gpu
				def     snake_case     (			self		:   Tuple							)     ->				List[Any]:




										'''simple docstring'''



										_UpperCamelCase      : Optional[int]         =   self._run_finetune(gpus=1							)
										self.assertGreaterEqual(result['''test'''][0]['''test_avg_em'''],			0.2							)






				@require_torch_multi_gpu
				def     snake_case     (			self		:   List[Any]							)     ->				Optional[int]:




										'''simple docstring'''



										_UpperCamelCase      : Optional[int]         =   self._run_finetune(gpus=2							)
										self.assertGreaterEqual(result['''test'''][0]['''test_avg_em'''],			0.2							)






				@require_torch_gpu
				@require_ray
				def     snake_case     (			self		:   List[str]							)     ->				Any:




										'''simple docstring'''



										_UpperCamelCase      : Union[str, Any]         =   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		:   int							)     ->				Union[str, Any]:




										'''simple docstring'''



										_UpperCamelCase      : Any         =   self._run_finetune(gpus=1,			distributed_retriever='''ray'''							)
										self.assertGreaterEqual(result['''test'''][0]['''test_avg_em'''],			0.2							)










"""simple docstring"""




import numpy as np
import torch
from imwatermark import WatermarkEncoder


# Copied from https://github.com/Stability-AI/generative-models/blob/613af104c6b85184091d42d374fef420eddb356d/scripts/demo/streamlit_helpers.py#L66
UpperCamelCase_									=0b1_0_1_1_0_0_1_1_1_1_1_0_1_1_0_0_1_0_0_1_0_0_0_0_0_1_1_1_1_0_1_1_1_0_1_1_0_0_0_1_1_0_0_1_1_1_1_0
# bin(x)[2:] gives bits of x as str, use int to convert them to 0/1
UpperCamelCase_									=[int(bit) for bit in bin(WATERMARK_MESSAGE)[2:]]




class 			_a						:






				def __init__(			self		:   str							)     ->				str:




										'''simple docstring'''



										_UpperCamelCase      : str         =   WATERMARK_BITS
										_UpperCamelCase      : Optional[int]         =   WatermarkEncoder()

										self.encoder.set_watermark('''bits''',			self.watermark							)






				def     snake_case     (			self		:   Dict,			lowerCAmelCase__		:   torch.FloatTensor							)     ->				int:




										'''simple docstring'''



										if images.shape[-1] < 2_5_6:
																return images

										_UpperCamelCase      : Union[str, Any]         =   (2_5_5 * (images / 2 + 0.5)).cpu().permute(0,			2,			3,			1							).float().numpy()

										_UpperCamelCase      : List[str]         =   [self.encoder.encode(lowerCAmelCase__,			'''dwtDct'''							) for image in images]

										_UpperCamelCase      : Dict         =   torch.from_numpy(np.array(lowerCAmelCase__							)							).permute(0,			3,			1,			2							)

										_UpperCamelCase      : Optional[int]         =   torch.clamp(2 * (images / 2_5_5 - 0.5),			min=-1.0,			max=1.0							)
										return images



