
'''simple docstring'''


import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


__lowercase							:   Dict					=     logging.get_logger(__name__)




def lowercase_	(   _lowercase			,       _lowercase			,       _lowercase			,       _lowercase=False				)     ->							Tuple:
					'''simple docstring'''

					try:
										import torch  # noqa: F401
					except ImportError:
										logger.error(
										    '''Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see'''
										    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
										    ''' instructions.'''				)
										raise

					if not is_sharded:
										lowerCamelCase_					:      str     =		os.path.abspath(_lowercase				)
										logger.info(F"""Loading PyTorch weights from {pt_path}"""				)

										lowerCamelCase_					:      Dict     =		torch.load(_lowercase			,       map_location='''cpu'''				)
										logger.info(F"""PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()				):,} parameters."""				)

										lowerCamelCase_					:      List[Any]     =		convert_pytorch_state_dict_to_flax(_lowercase			,       _lowercase				)
					else:
										# model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
										lowerCamelCase_					:      Optional[int]     =		convert_pytorch_sharded_state_dict_to_flax(_lowercase			,       _lowercase				)
					return flax_state_dict




def lowercase_	(   _lowercase			,       _lowercase			,       _lowercase			,       _lowercase			,       )     ->							(Tuple[str], np.ndarray):
					'''simple docstring'''

					def is_key_or_prefix_key_in_dict(_lowercase				) -> bool:
										return len(set(_lowercase				) & {key, (model_prefix,) + key}				) > 0

					# layer norm
					lowerCamelCase_					:      Union[str, Any]     =		pt_tuple_key[:-1] + ('''scale''',)
					if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(_lowercase				):
										return renamed_pt_tuple_key, pt_tensor

					# batch norm layer mean
					lowerCamelCase_					:      Optional[Any]     =		pt_tuple_key[:-1] + ('''mean''',)
					if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(_lowercase				):
										return renamed_pt_tuple_key, pt_tensor

					# batch norm layer var
					lowerCamelCase_					:      Any     =		pt_tuple_key[:-1] + ('''var''',)
					if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(_lowercase				):
										return renamed_pt_tuple_key, pt_tensor

					# embedding
					lowerCamelCase_					:      Union[str, Any]     =		pt_tuple_key[:-1] + ('''embedding''',)
					if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(_lowercase				):
										return renamed_pt_tuple_key, pt_tensor

					# conv layer
					lowerCamelCase_					:      Union[str, Any]     =		pt_tuple_key[:-1] + ('''kernel''',)
					if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(_lowercase				):
										lowerCamelCase_					:      Union[str, Any]     =		pt_tensor.transpose(2			,       3			,       1			,       0				)
										return renamed_pt_tuple_key, pt_tensor

					# linear layer
					lowerCamelCase_					:      int     =		pt_tuple_key[:-1] + ('''kernel''',)
					if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(_lowercase				):
										lowerCamelCase_					:      List[str]     =		pt_tensor.T
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm weight
					lowerCamelCase_					:      Optional[Any]     =		pt_tuple_key[:-1] + ('''weight''',)
					if pt_tuple_key[-1] == "gamma":
										return renamed_pt_tuple_key, pt_tensor

					# old PyTorch layer norm bias
					lowerCamelCase_					:      Any     =		pt_tuple_key[:-1] + ('''bias''',)
					if pt_tuple_key[-1] == "beta":
										return renamed_pt_tuple_key, pt_tensor

					# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
					lowerCamelCase_					:      Tuple     =		None
					if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
										lowerCamelCase_					:      Optional[int]     =		pt_tuple_key[-2] + '''_g'''
					elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
										lowerCamelCase_					:      Dict     =		pt_tuple_key[-2] + '''_v'''
					if name is not None:
										lowerCamelCase_					:      Dict     =		pt_tuple_key[:-3] + (name,)
										return renamed_pt_tuple_key, pt_tensor

					return pt_tuple_key, pt_tensor




def lowercase_	(   _lowercase			,       _lowercase				)     ->							Tuple:
					'''simple docstring'''

					lowerCamelCase_					:      Union[str, Any]     =		{k: v.numpy() for k, v in pt_state_dict.items()}

					lowerCamelCase_					:      List[str]     =		flax_model.base_model_prefix

					# use params dict if the model contains batch norm layers
					if "params" in flax_model.params:
										lowerCamelCase_					:      Any     =		flax_model.params['''params''']
					else:
										lowerCamelCase_					:      Tuple     =		flax_model.params
					lowerCamelCase_					:      int     =		flatten_dict(_lowercase				)

					# add batch_stats keys,values to dict
					if "batch_stats" in flax_model.params:
										lowerCamelCase_					:      int     =		flatten_dict(flax_model.params['''batch_stats''']				)
										random_flax_state_dict.update(_lowercase				)

					lowerCamelCase_					:      Any     =		{}

					lowerCamelCase_					:      Any     =		(model_prefix not in flax_model_params) and (
					    model_prefix in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
					)
					lowerCamelCase_					:      int     =		(model_prefix in flax_model_params) and (
					    model_prefix not in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
					)

					# Need to change some parameters name to match Flax names
					for pt_key, pt_tensor in pt_state_dict.items():
										lowerCamelCase_					:      int     =		tuple(pt_key.split('''.'''				)				)

										# remove base model prefix if necessary
										lowerCamelCase_					:      Optional[int]     =		pt_tuple_key[0] == model_prefix
										if load_model_with_head_into_base_model and has_base_model_prefix:
															lowerCamelCase_					:      List[str]     =		pt_tuple_key[1:]

										# Correctly rename weight parameters
										lowerCamelCase_					:      str     =		rename_key_and_reshape_tensor(
										    _lowercase			,       _lowercase			,       _lowercase			,       _lowercase				)

										# add model prefix if necessary
										lowerCamelCase_					:      Tuple     =		(model_prefix,) + flax_key in random_flax_state_dict
										if load_base_model_into_model_with_head and require_base_model_prefix:
															lowerCamelCase_					:      Optional[int]     =		(model_prefix,) + flax_key

										if flax_key in random_flax_state_dict:
															if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																				raise ValueError(
																				    F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
																				    F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""				)

        # add batch stats if the model contains batchnorm layers
										if "batch_stats" in flax_model.params:
															if "mean" in flax_key[-1] or "var" in flax_key[-1]:
																				lowerCamelCase_					:      Any     =		jnp.asarray(_lowercase				)
																				continue
															# remove num_batches_tracked key
															if "num_batches_tracked" in flax_key[-1]:
																				flax_state_dict.pop(_lowercase			,       _lowercase				)
																				continue

															# also add unexpected weight so that warning is thrown
															lowerCamelCase_					:      List[str]     =		jnp.asarray(_lowercase				)

										else:
															# also add unexpected weight so that warning is thrown
															lowerCamelCase_					:      Union[str, Any]     =		jnp.asarray(_lowercase				)

					return unflatten_dict(_lowercase				)




def lowercase_	(   _lowercase			,       _lowercase				)     ->							Optional[Any]:
					'''simple docstring'''

					import torch

					# Load the index
					lowerCamelCase_					:      Dict     =		{}
					for shard_file in shard_filenames:
										# load using msgpack utils
										lowerCamelCase_					:      Optional[Any]     =		torch.load(_lowercase				)
										lowerCamelCase_					:      Any     =		{k: v.numpy() for k, v in pt_state_dict.items()}

										lowerCamelCase_					:      Optional[Any]     =		flax_model.base_model_prefix

										# use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
										if "batch_stats" in flax_model.params:
															lowerCamelCase_					:      List[Any]     =		flax_model.params['''params''']

															lowerCamelCase_					:      Union[str, Any]     =		flatten_dict(_lowercase				)
															random_flax_state_dict.update(flatten_dict(flax_model.params['''batch_stats''']				)				)
										else:
															lowerCamelCase_					:      List[Any]     =		flax_model.params
															lowerCamelCase_					:      Union[str, Any]     =		flatten_dict(_lowercase				)

										lowerCamelCase_					:      Tuple     =		(model_prefix not in flax_model_params) and (
										    model_prefix in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
										)
										lowerCamelCase_					:      int     =		(model_prefix in flax_model_params) and (
										    model_prefix not in {k.split('''.'''				)[0] for k in pt_state_dict.keys()}
										)
										# Need to change some parameters name to match Flax names
										for pt_key, pt_tensor in pt_state_dict.items():
															lowerCamelCase_					:      Dict     =		tuple(pt_key.split('''.'''				)				)

															# remove base model prefix if necessary
															lowerCamelCase_					:      Optional[Any]     =		pt_tuple_key[0] == model_prefix
															if load_model_with_head_into_base_model and has_base_model_prefix:
																				lowerCamelCase_					:      Dict     =		pt_tuple_key[1:]

															# Correctly rename weight parameters
															lowerCamelCase_					:      Optional[int]     =		rename_key_and_reshape_tensor(
															    _lowercase			,       _lowercase			,       _lowercase			,       _lowercase				)
															# add model prefix if necessary
															lowerCamelCase_					:      Optional[Any]     =		(model_prefix,) + flax_key in random_flax_state_dict
															if load_base_model_into_model_with_head and require_base_model_prefix:
																				lowerCamelCase_					:      Optional[int]     =		(model_prefix,) + flax_key

															if flax_key in random_flax_state_dict:
																				if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
																									raise ValueError(
																									    F"""PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape """
																									    F"""{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."""				)

            # add batch stats if the model contains batchnorm layers
															if "batch_stats" in flax_model.params:
																				if "mean" in flax_key[-1]:
																									lowerCamelCase_					:      List[str]     =		jnp.asarray(_lowercase				)
																									continue
																				if "var" in flax_key[-1]:
																									lowerCamelCase_					:      Tuple     =		jnp.asarray(_lowercase				)
																									continue
																				# remove num_batches_tracked key
																				if "num_batches_tracked" in flax_key[-1]:
																									flax_state_dict.pop(_lowercase			,       _lowercase				)
																									continue

																				# also add unexpected weight so that warning is thrown
																				lowerCamelCase_					:      Optional[Any]     =		jnp.asarray(_lowercase				)

															else:
																				# also add unexpected weight so that warning is thrown
																				lowerCamelCase_					:      List[Any]     =		jnp.asarray(_lowercase				)
					return unflatten_dict(_lowercase				)




def lowercase_	(   _lowercase			,       _lowercase				)     ->							List[Any]:
					'''simple docstring'''

					lowerCamelCase_					:      List[Any]     =		os.path.abspath(_lowercase				)
					logger.info(F"""Loading Flax weights from {flax_checkpoint_path}"""				)

					# import correct flax class
					lowerCamelCase_					:      Dict     =		getattr(_lowercase			,       '''Flax''' + model.__class__.__name__				)

					# load flax weight dict
					with open(_lowercase			,       '''rb'''				) as state_f:
										try:
															lowerCamelCase_					:      Optional[Any]     =		from_bytes(_lowercase			,       state_f.read()				)
										except UnpicklingError:
															raise EnvironmentError(F"""Unable to convert {flax_checkpoint_path} to Flax deserializable object. """				)

					return load_flax_weights_in_pytorch_model(_lowercase			,       _lowercase				)






def lowercase_	(   _lowercase			,       _lowercase				)     ->							Tuple:
					'''simple docstring'''

					try:
										import torch  # noqa: F401
					except ImportError:
										logger.error(
										    '''Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see'''
										    ''' https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation'''
										    ''' instructions.'''				)
										raise

					# check if we have bf16 weights
					lowerCamelCase_					:      List[Any]     =		flatten_dict(jax.tree_util.tree_map(lambda _lowercase				: x.dtype == jnp.bfloataa			,       _lowercase				)				).values()
					if any(_lowercase				):
										# convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
										# and bf16 is not fully supported in PT yet.
										logger.warning(
										    '''Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` '''
										    '''before loading those in PyTorch model.'''				)
										lowerCamelCase_					:      Tuple     =		jax.tree_util.tree_map(
										    lambda _lowercase				: params.astype(np.floataa				) if params.dtype == jnp.bfloataa else params			,       _lowercase				)

					lowerCamelCase_					:      Any     =		flatten_dict(_lowercase				)
					lowerCamelCase_					:      List[Any]     =		pt_model.state_dict()

					lowerCamelCase_					:      Union[str, Any]     =		(pt_model.base_model_prefix in flax_state) and (
					    pt_model.base_model_prefix not in {k.split('''.'''				)[0] for k in pt_model_dict.keys()}
					)
					lowerCamelCase_					:      Optional[Any]     =		(pt_model.base_model_prefix not in flax_state) and (
					    pt_model.base_model_prefix in {k.split('''.'''				)[0] for k in pt_model_dict.keys()}
					)

					# keep track of unexpected & missing keys
					lowerCamelCase_					:      List[Any]     =		[]
					lowerCamelCase_					:      Optional[int]     =		set(pt_model_dict.keys()				)

					for flax_key_tuple, flax_tensor in flax_state_dict.items():
										lowerCamelCase_					:      List[str]     =		flax_key_tuple[0] == pt_model.base_model_prefix
										lowerCamelCase_					:      List[str]     =		'''.'''.join((pt_model.base_model_prefix,) + flax_key_tuple				) in pt_model_dict

										# adapt flax_key to prepare for loading from/to base model only
										if load_model_with_head_into_base_model and has_base_model_prefix:
															lowerCamelCase_					:      List[str]     =		flax_key_tuple[1:]
										elif load_base_model_into_model_with_head and require_base_model_prefix:
															lowerCamelCase_					:      Optional[Any]     =		(pt_model.base_model_prefix,) + flax_key_tuple

										# rename flax weights to PyTorch format
										if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(_lowercase				) not in pt_model_dict:
															# conv layer
															lowerCamelCase_					:      Any     =		flax_key_tuple[:-1] + ('''weight''',)
															lowerCamelCase_					:      Union[str, Any]     =		jnp.transpose(_lowercase			,       (3, 2, 0, 1)				)
										elif flax_key_tuple[-1] == "kernel" and ".".join(_lowercase				) not in pt_model_dict:
															# linear layer
															lowerCamelCase_					:      str     =		flax_key_tuple[:-1] + ('''weight''',)
															lowerCamelCase_					:      Union[str, Any]     =		flax_tensor.T
										elif flax_key_tuple[-1] in ["scale", "embedding"]:
															lowerCamelCase_					:      int     =		flax_key_tuple[:-1] + ('''weight''',)

										# adding batch stats from flax batch norm to pt
										elif "mean" in flax_key_tuple[-1]:
															lowerCamelCase_					:      Optional[int]     =		flax_key_tuple[:-1] + ('''running_mean''',)
										elif "var" in flax_key_tuple[-1]:
															lowerCamelCase_					:      Dict     =		flax_key_tuple[:-1] + ('''running_var''',)

										if "batch_stats" in flax_state:
															lowerCamelCase_					:      List[Any]     =		'''.'''.join(flax_key_tuple[1:]				)  # Remove the params/batch_stats header
										else:
															lowerCamelCase_					:      Optional[Any]     =		'''.'''.join(_lowercase				)

										# We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
										lowerCamelCase_					:      List[str]     =		{}
										# New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
										for key in pt_model_dict:
															lowerCamelCase_					:      List[Any]     =		key.split('''.'''				)
															lowerCamelCase_					:      int     =		None
															if key_components[-3::2] == ["parametrizations", "original0"]:
																				lowerCamelCase_					:      Dict     =		key_components[-2] + '''_g'''
															elif key_components[-3::2] == ["parametrizations", "original1"]:
																				lowerCamelCase_					:      Any     =		key_components[-2] + '''_v'''
															if name is not None:
																				lowerCamelCase_					:      Union[str, Any]     =		key_components[:-3] + [name]
																				lowerCamelCase_					:      List[str]     =		'''.'''.join(_lowercase				)
																				lowerCamelCase_					:      Any     =		key

										if flax_key in special_pt_names:
															lowerCamelCase_					:      List[str]     =		special_pt_names[flax_key]

										if flax_key in pt_model_dict:
															if flax_tensor.shape != pt_model_dict[flax_key].shape:
																				raise ValueError(
																				    F"""Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected """
																				    F"""to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}."""				)
															else:
																				# add weight to pytorch dict
																				lowerCamelCase_					:      Optional[int]     =		np.asarray(_lowercase				) if not isinstance(_lowercase			,       np.ndarray				) else flax_tensor
																				lowerCamelCase_					:      Optional[Any]     =		torch.from_numpy(_lowercase				)
																				# remove from missing keys
																				missing_keys.remove(_lowercase				)
										else:
															# weight is not expected by PyTorch model
															unexpected_keys.append(_lowercase				)

					pt_model.load_state_dict(_lowercase				)

					# re-transform missing_keys to list
					lowerCamelCase_					:      Optional[Any]     =		list(_lowercase				)

					if len(_lowercase				) > 0:
										logger.warning(
										    '''Some weights of the Flax model were not used when initializing the PyTorch model'''
										    F""" {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing"""
										    F""" {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture"""
										    ''' (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This'''
										    F""" IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect"""
										    ''' to be exactly identical (e.g. initializing a BertForSequenceClassification model from a'''
										    ''' FlaxBertForSequenceClassification model).'''				)
					else:
										logger.warning(F"""All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n"""				)
					if len(_lowercase				) > 0:
										logger.warning(
										    F"""Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly"""
										    F""" initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to"""
										    ''' use it for predictions and inference.'''				)
					else:
										logger.warning(
										    F"""All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n"""
										    '''If your task is similar to the task the model of the checkpoint was trained on, '''
										    F"""you can already use {pt_model.__class__.__name__} for predictions without further training."""				)

					return pt_model


'''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)
__lowercase							:   Tuple					=     logging.getLogger()

__lowercase							:   str					=     logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)

class 	__lowercase				(  _lowercase    ):

			def        UpperCAmelCase__ (self			,   A      ):
								os.makedirs(A			,   exist_ok=A      )
								lowerCamelCase_					:      Any     =		{'''source''': '''What is love ?''', '''target''': '''life'''}
								lowerCamelCase_					:      List[str]     =		{'''train''': 1_2, '''val''': 2, '''test''': 2}
								for split in ["train", "test", "val"]:
													for field in ["source", "target"]:
																		lowerCamelCase_					:      List[Any]     =		'''\n'''.join([contents[field]] * n_lines[split]      )
																		with open(os.path.join(A			,   F"""{split}.{field}"""      )			,   '''w'''      ) as f:
																							f.write(A      )

			def        UpperCAmelCase__ (self			,   A			,   A = "pytorch"      ):
								lowerCamelCase_					:      List[str]     =		self.get_auto_remove_tmp_dir()
								lowerCamelCase_					:      Tuple     =		os.path.join(A			,   '''output'''      )
								lowerCamelCase_					:      Union[str, Any]     =		os.path.join(A			,   '''data'''      )
								self._create_dummy_data(data_dir=A      )

								lowerCamelCase_					:      Tuple     =		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'''      )

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

								lowerCamelCase_					:      List[Any]     =		os.path.join(A			,   '''metrics.json'''      )
								with open(A      ) as f:
													lowerCamelCase_					:      List[str]     =		json.load(A      )
								return result

			@require_torch_gpu
			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Tuple     =		self._run_finetune(gpus=1      )
								self.assertGreaterEqual(result['''test'''][0]['''test_avg_em''']			,   0.2      )

			@require_torch_multi_gpu
			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      List[str]     =		self._run_finetune(gpus=2      )
								self.assertGreaterEqual(result['''test'''][0]['''test_avg_em''']			,   0.2      )

			@require_torch_gpu
			@require_ray
			def        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Optional[int]     =		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        UpperCAmelCase__ (self      ):
								lowerCamelCase_					:      Union[str, Any]     =		self._run_finetune(gpus=1			,   distributed_retriever='''ray'''      )
								self.assertGreaterEqual(result['''test'''][0]['''test_avg_em''']			,   0.2      )


