

import unittest
from unittest.util import safe_repr

from transformers import AutoTokenizer, RwkvConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    RWKV_PRETRAINED_MODEL_ARCHIVE_LIST,
							    RwkvForCausalLM,
							    RwkvModel,
							)
							from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_0
else:
							lowerCAmelCase__             =		False







class 						lowercase     :


		"""simple docstring"""




		def __init__(					self						,							__snake_case						,							__snake_case=14						,							__snake_case=7						,							__snake_case=True						,							__snake_case=False						,							__snake_case=True						,							__snake_case=True						,							__snake_case=True						,							__snake_case=99						,							__snake_case=32						,							__snake_case=5						,							__snake_case=37						,							__snake_case="gelu"						,							__snake_case=0.1						,							__snake_case=0.1						,							__snake_case=5_12						,							__snake_case=16						,							__snake_case=2						,							__snake_case=3						,							__snake_case=4						,							__snake_case=None						,							):
			_UpperCamelCase :		Tuple								=    parent
			_UpperCamelCase :		Tuple								=    batch_size
			_UpperCamelCase :		List[str]								=    seq_length
			_UpperCamelCase :		Union[str, Any]								=    is_training
			_UpperCamelCase :		Union[str, Any]								=    use_token_type_ids
			_UpperCamelCase :		int								=    use_input_mask
			_UpperCamelCase :		Optional[int]								=    use_labels
			_UpperCamelCase :		str								=    use_mc_token_ids
			_UpperCamelCase :		str								=    vocab_size
			_UpperCamelCase :		int								=    hidden_size
			_UpperCamelCase :		Dict								=    num_hidden_layers
			_UpperCamelCase :		str								=    intermediate_size
			_UpperCamelCase :		str								=    hidden_act
			_UpperCamelCase :		Union[str, Any]								=    hidden_dropout_prob
			_UpperCamelCase :		str								=    attention_probs_dropout_prob
			_UpperCamelCase :		Tuple								=    max_position_embeddings
			_UpperCamelCase :		Any								=    type_vocab_size
			_UpperCamelCase :		List[str]								=    type_sequence_label_size
			_UpperCamelCase :		str								=    num_labels
			_UpperCamelCase :		List[Any]								=    num_choices
			_UpperCamelCase :		Any								=    scope
			_UpperCamelCase :		Dict								=    vocab_size - 1
			_UpperCamelCase :		Tuple								=    vocab_size - 1
			_UpperCamelCase :		Any								=    vocab_size - 1


		def 		A__       (					self):
			return RwkvConfig.from_pretrained('sgugger/rwkv-4-pile-7b')


		def 		A__       (					self						,							__snake_case=False						,							__snake_case=False						,							__snake_case=False):
			_UpperCamelCase :		int								=    ids_tensor([self.batch_size, self.seq_length]						,							self.vocab_size)

			_UpperCamelCase :		Union[str, Any]								=    None
			if self.use_input_mask:
				_UpperCamelCase :		str								=    random_attention_mask([self.batch_size, self.seq_length])

			_UpperCamelCase :		Union[str, Any]								=    None
			if self.use_token_type_ids:
				_UpperCamelCase :		Optional[int]								=    ids_tensor([self.batch_size, self.seq_length]						,							self.type_vocab_size)

			_UpperCamelCase :		Any								=    None
			if self.use_mc_token_ids:
				_UpperCamelCase :		Union[str, Any]								=    ids_tensor([self.batch_size, self.num_choices]						,							self.seq_length)

			_UpperCamelCase :		Union[str, Any]								=    None
			_UpperCamelCase :		str								=    None
			_UpperCamelCase :		str								=    None
			if self.use_labels:
				_UpperCamelCase :		Tuple								=    ids_tensor([self.batch_size]						,							self.type_sequence_label_size)
				_UpperCamelCase :		Optional[Any]								=    ids_tensor([self.batch_size, self.seq_length]						,							self.num_labels)
				_UpperCamelCase :		Union[str, Any]								=    ids_tensor([self.batch_size]						,							self.num_choices)

			_UpperCamelCase :		Optional[int]								=    self.get_config(
			    gradient_checkpointing=__snake_case						,							scale_attn_by_inverse_layer_idx=__snake_case						,							reorder_and_upcast_attn=__snake_case						,							)

			return (
			    config,
			    input_ids,
			    input_mask,
			    None,
			    token_type_ids,
			    mc_token_ids,
			    sequence_labels,
			    token_labels,
			    choice_labels,
			)


		def 		A__       (					self						,							__snake_case=False						,							__snake_case=False						,							__snake_case=False):
			return RwkvConfig(
			    vocab_size=self.vocab_size						,							hidden_size=self.hidden_size						,							num_hidden_layers=self.num_hidden_layers						,							intermediate_size=self.intermediate_size						,							activation_function=self.hidden_act						,							resid_pdrop=self.hidden_dropout_prob						,							attn_pdrop=self.attention_probs_dropout_prob						,							n_positions=self.max_position_embeddings						,							type_vocab_size=self.type_vocab_size						,							use_cache=__snake_case						,							bos_token_id=self.bos_token_id						,							eos_token_id=self.eos_token_id						,							pad_token_id=self.pad_token_id						,							gradient_checkpointing=__snake_case						,							scale_attn_by_inverse_layer_idx=__snake_case						,							reorder_and_upcast_attn=__snake_case						,							)


		def 		A__       (					self):
			_UpperCamelCase :		Any								=    self.get_config()
			_UpperCamelCase :		int								=    3_00
			return config


		def 		A__       (					self):
			(
			    (
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		
			) :		List[str]								=    self.prepare_config_and_inputs()

			_UpperCamelCase :		Optional[int]								=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
			_UpperCamelCase :		List[Any]								=    ids_tensor([self.batch_size, self.seq_length]						,							vocab_size=2)

			return (
			    config,
			    input_ids,
			    input_mask,
			    head_mask,
			    token_type_ids,
			    sequence_labels,
			    token_labels,
			    choice_labels,
			    encoder_hidden_states,
			    encoder_attention_mask,
			)


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							*__snake_case):
			_UpperCamelCase :		Optional[Any]								=    True
			_UpperCamelCase :		Tuple								=    RwkvModel(config=__snake_case)
			model.to(__snake_case)
			model.eval()

			_UpperCamelCase :		Optional[int]								=    model(__snake_case)

			self.parent.assertEqual(result.last_hidden_state.shape						,							(self.batch_size, self.seq_length, self.hidden_size))
			self.parent.assertEqual(len(result.hidden_states)						,							config.num_hidden_layers + 1)


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							*__snake_case):
			_UpperCamelCase :		Dict								=    RwkvForCausalLM(__snake_case)
			model.to(__snake_case)
			model.eval()

			_UpperCamelCase :		int								=    model(__snake_case						,							labels=__snake_case)
			self.parent.assertEqual(result.loss.shape						,							())
			self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.seq_length, self.vocab_size))


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							*__snake_case):
			_UpperCamelCase :		List[Any]								=    RwkvModel(config=__snake_case)
			model.to(__snake_case)
			model.eval()

			_UpperCamelCase :		List[str]								=    model(__snake_case)
			_UpperCamelCase :		Any								=    outputs.last_hidden_state

			_UpperCamelCase :		Dict								=    model(input_ids[:, :2])
			_UpperCamelCase :		List[Any]								=    outputs.last_hidden_state

			# Using the state computed on the first inputs, we will get the same output
			_UpperCamelCase :		List[str]								=    model(input_ids[:, 2:]						,							state=outputs.state)
			_UpperCamelCase :		Optional[Any]								=    outputs.last_hidden_state

			self.parent.assertTrue(torch.allclose(torch.cat([output_one, output_two]						,							dim=1)						,							__snake_case						,							atol=1e-5))


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							__snake_case						,							*__snake_case						,							__snake_case=False):
			_UpperCamelCase :		Optional[int]								=    RwkvForCausalLM(__snake_case)
			model.to(__snake_case)
			if gradient_checkpointing:
				model.gradient_checkpointing_enable()

			_UpperCamelCase :		Any								=    model(__snake_case						,							labels=__snake_case)
			self.parent.assertEqual(result.loss.shape						,							())
			self.parent.assertEqual(result.logits.shape						,							(self.batch_size, self.seq_length, self.vocab_size))
			result.loss.backward()






		def 		A__       (					self):
			_UpperCamelCase :		List[str]								=    self.prepare_config_and_inputs()

			(
			    (
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		(
			    _UpperCamelCase
			)						,		
			) :		Optional[int]								=    config_and_inputs

			_UpperCamelCase :		Dict								=    {'input_ids': input_ids}

			return config, inputs_dict







@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0     ,				reason="See https://github.com/huggingface/transformers/pull/24204"				)
@require_torch
class 						lowercase     (					_lowercase     ,				_lowercase     ,				_lowercase     ,				unittest.TestCase				):


		"""simple docstring"""




		a__       							=	(RwkvModel, RwkvForCausalLM) if is_torch_available() else ()
		a__       							=	(
		    {"feature-extraction": RwkvModel, "text-generation": RwkvForCausalLM} if is_torch_available() else {}
		)
		# all_generative_model_classes = (RwkvForCausalLM,) if is_torch_available() else ()
		a__       							=	False
		a__       							=	False
		a__       							=	False
		a__       							=	False
		a__       							=	False  # Rwkv does not support head masking


		def 		A__       (					self):
			_UpperCamelCase :		int								=    RwkvModelTester(self)
			_UpperCamelCase :		Tuple								=    ConfigTester(
			    self						,							config_class=__snake_case						,							n_embd=37						,							common_properties=['hidden_size', 'num_hidden_layers'])


		def 		A__       (					self						,							__snake_case						,							__snake_case						,							__snake_case=None):
			if isinstance(__snake_case						,							torch.Tensor):
				_UpperCamelCase						,		_UpperCamelCase :		Any								=    member.max().item(), member.min().item()
			elif isinstance(__snake_case						,							__snake_case) or isinstance(__snake_case						,							__snake_case):
				_UpperCamelCase						,		_UpperCamelCase :		Union[str, Any]								=    max(__snake_case), min(__snake_case)

			if not isinstance(__snake_case						,							__snake_case):
				raise TypeError('container should be a list or tuple')
			elif len(__snake_case) != 2:
				raise ValueError('container should have 2 elements')

			_UpperCamelCase						,		_UpperCamelCase :		Optional[int]								=    container

			_UpperCamelCase :		str								=    (min_value >= expected_min) and (max_value <= expected_max)

			if not is_inside_interval:
				_UpperCamelCase :		Any								=    '%s not found in %s' % (safe_repr(__snake_case), safe_repr(__snake_case))
				self.fail(self._formatMessage(__snake_case						,							__snake_case))


		def 		A__       (					self):
			self.config_tester.run_common_tests()


		def 		A__       (					self):
			_UpperCamelCase :		Tuple								=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_rwkv_model(*__snake_case)


		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_causl_lm(*__snake_case)


		def 		A__       (					self):
			_UpperCamelCase :		Optional[int]								=    self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_state_equivalency(*__snake_case)


		def 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		List[str]								=    self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
				_UpperCamelCase :		Optional[Any]								=    model_class(config=__snake_case)
				for name, param in model.named_parameters():
					if "time_decay" in name:
						if param.requires_grad:
							self.assertTrue(param.data.max().item() == 3.0)
							self.assertTrue(param.data.min().item() == -5.0)
					elif "time_first" in name:
						if param.requires_grad:
							# check if it's a ones like
							self.assertTrue(torch.allclose(param.data						,							torch.ones_like(param.data)						,							atol=1e-5						,							rtol=1e-5))
					elif any(x in name for x in ['time_mix_key', 'time_mix_receptance']):
						if param.requires_grad:
							self.assertInterval(
							    param.data						,							[0.0, 1.0]						,							msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''						,							)
					elif "time_mix_value" in name:
						if param.requires_grad:
							self.assertInterval(
							    param.data						,							[0.0, 1.3]						,							msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''						,							)


		def 		A__       (					self):
			_UpperCamelCase						,		_UpperCamelCase :		List[Any]								=    self.model_tester.prepare_config_and_inputs_for_common()
			_UpperCamelCase :		Optional[int]								=    True

			_UpperCamelCase :		Optional[int]								=    getattr(self.model_tester						,							'seq_length'						,							__snake_case)

			for model_class in self.all_model_classes:
				_UpperCamelCase :		Optional[Any]								=    True
				_UpperCamelCase :		Any								=    False
				_UpperCamelCase :		Optional[int]								=    True
				_UpperCamelCase :		str								=    model_class(__snake_case)
				model.to(__snake_case)
				model.eval()

				_UpperCamelCase :		List[str]								=    self._prepare_for_class(__snake_case						,							__snake_case)
				_UpperCamelCase :		Dict								=    inputs['input_ids'].shape[0]
				with torch.no_grad():
					_UpperCamelCase :		int								=    model(**__snake_case)
				_UpperCamelCase :		Optional[Any]								=    outputs.attentions
				self.assertEqual(len(__snake_case)						,							self.model_tester.num_hidden_layers)

				# check that output_attentions also work using config
				del inputs_dict["output_attentions"]
				_UpperCamelCase :		Optional[Any]								=    True
				_UpperCamelCase :		int								=    model_class(__snake_case)
				model.to(__snake_case)
				model.eval()

				_UpperCamelCase :		Dict								=    self._prepare_for_class(__snake_case						,							__snake_case)
				_UpperCamelCase :		Optional[Any]								=    inputs['input_ids'].shape[0]
				with torch.no_grad():
					_UpperCamelCase :		Optional[Any]								=    model(**__snake_case)
				_UpperCamelCase :		int								=    outputs.attentions
				self.assertEqual(len(__snake_case)						,							self.model_tester.num_hidden_layers)

				self.assertListEqual(
				    list(attentions[0].shape[-3:])						,							[batch_size, seq_len, config.hidden_size]						,							)
				_UpperCamelCase :		Any								=    len(__snake_case)

				# Check attention is always last and order is fine
				_UpperCamelCase :		Any								=    True
				_UpperCamelCase :		Union[str, Any]								=    True
				_UpperCamelCase :		int								=    model_class(__snake_case)
				model.to(__snake_case)
				model.eval()

				_UpperCamelCase :		Union[str, Any]								=    self._prepare_for_class(__snake_case						,							__snake_case)
				_UpperCamelCase :		Optional[int]								=    inputs['input_ids'].shape[0]
				with torch.no_grad():
					_UpperCamelCase :		Dict								=    model(**__snake_case)

				_UpperCamelCase :		Any								=    1
				self.assertEqual(out_len + added_hidden_states						,							len(__snake_case))

				_UpperCamelCase :		Optional[int]								=    outputs.attentions

				self.assertEqual(len(__snake_case)						,							self.model_tester.num_hidden_layers)
				self.assertListEqual(
				    list(self_attentions[0].shape[-3:])						,							[batch_size, seq_len, config.hidden_size]						,							)






		@slow
		def 		A__       (					self):
			for model_name in RWKV_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
				_UpperCamelCase :		str								=    RwkvModel.from_pretrained(__snake_case)
				self.assertIsNotNone(__snake_case)







@unittest.skipIf(
    not is_torch_greater_or_equal_than_2_0     ,				reason="See https://github.com/huggingface/transformers/pull/24204"				)
@slow
class 						lowercase     (					unittest.TestCase				):


		"""simple docstring"""




		def 		A__       (					self):
			_UpperCamelCase :		Optional[int]								=    'RWKV/rwkv-4-169m-pile'
			_UpperCamelCase :		Any								=    AutoTokenizer.from_pretrained(self.model_id)


		def 		A__       (					self):
			_UpperCamelCase :		int								=    'Hello my name is Jasmine and I am a newbie to the'
			_UpperCamelCase :		Optional[Any]								=    RwkvForCausalLM.from_pretrained(self.model_id).to(__snake_case)

			_UpperCamelCase :		Optional[Any]								=    self.tokenizer('Hello my name is'						,							return_tensors='pt').input_ids.to(__snake_case)
			_UpperCamelCase :		str								=    model.generate(__snake_case						,							max_new_tokens=10)
			_UpperCamelCase :		Any								=    self.tokenizer.decode(output[0].tolist())

			self.assertEqual(__snake_case						,							__snake_case)






		def 		A__       (					self):
			_UpperCamelCase :		Union[str, Any]								=    'Hello my name is Jasmine and I am a newbie to the'

			_UpperCamelCase :		int								=    self.tokenizer('Hello my name is'						,							return_tensors='pt').input_ids.to(__snake_case)
			_UpperCamelCase :		str								=    RwkvForCausalLM.from_pretrained(self.model_id						,							torch_dtype=torch.bfloataa).to(__snake_case)

			_UpperCamelCase :		Union[str, Any]								=    model.generate(__snake_case						,							max_new_tokens=10)
			_UpperCamelCase :		Union[str, Any]								=    self.tokenizer.decode(output[0].tolist())

			self.assertEqual(__snake_case						,							__snake_case)







import argparse
import itertools
import operator
import sys
from collections import OrderedDict

from run_eval import datetime_now, run_generate

from utils import ROUGE_KEYS


# A table of supported tasks and the list of scores in the order of importance to be sorted by.
# To add a new task, simply list the score names that `run_eval.run_generate()` returns
lowerCAmelCase__             =		{
    """translation""": ["""bleu"""],
    """summarization""": ROUGE_KEYS,
}






def        lowerCamelCase_  (				UpperCAmelCase_				:  Optional[int]						)    ->					Any:




	'''simple docstring'''




	_UpperCamelCase :		str								=    search.split()
	_UpperCamelCase :		str								=    dict((g.split('='						) for g in groups)						)
	_UpperCamelCase :		List[Any]								=    list(entries.keys()						)
	_UpperCamelCase :		Dict								=    [[F'''--{k} {v}''' for v in vs.split(':'						)] for k, vs in entries.items()]
	_UpperCamelCase :		Dict								=    [list(UpperCAmelCase_						) for x in itertools.product(*UpperCAmelCase_						)]
	return matrix, entry_names






def        lowerCamelCase_  (				)    ->					int:




	'''simple docstring'''




	_UpperCamelCase :		Optional[Any]								=    sys.argv[0]

	_UpperCamelCase :		Tuple								=    argparse.ArgumentParser(
	    usage=(
	        '\n\nImportant: this script accepts all arguments `run_eval.py` accepts and then a few extra, therefore'
	        ' refer to `run_eval.py -h` for the complete list.'
	    )						)
	parser.add_argument(
	    '--search'				,							type=UpperCAmelCase_				,							required=UpperCAmelCase_				,							help='param space to search, e.g. "num_beams=5:10 length_penalty=0.8:1.0:1.2"'				,							)
	parser.add_argument(
	    '--bs'				,							type=UpperCAmelCase_				,							default=8				,							required=UpperCAmelCase_				,							help='initial batch size (may get reduced if it\'s too big)'						)
	parser.add_argument('--task'				,							type=UpperCAmelCase_				,							help='used for task_specific_params + metrics'						)
	parser.add_argument(
	    '--info'				,							nargs='?'				,							type=UpperCAmelCase_				,							const=datetime_now()				,							help=(
	        'add custom notes to be printed before the results table. If no value is passed, the current datetime'
	        ' string will be used.'
	    )				,							)
	_UpperCamelCase						,		_UpperCamelCase :		str								=    parser.parse_known_args()
	# we share some of the args
	args_main.extend(['--task', args.task]						)
	_UpperCamelCase :		str								=    [prog] + args_main

	# to support variations like translation_en_to_de"
	_UpperCamelCase :		Tuple								=    'translation' if 'translation' in args.task else 'summarization'

	_UpperCamelCase						,		_UpperCamelCase :		Optional[int]								=    parse_search_arg(args.search						)
	_UpperCamelCase :		List[Any]								=    task_score_names[task]  # score cols first
	_UpperCamelCase :		Optional[int]								=    {col: len(str(UpperCAmelCase_						)						) for col in col_names}
	_UpperCamelCase :		str								=    []
	for r in matrix:
		_UpperCamelCase :		Union[str, Any]								=    dict((x.replace('--'				,							''						).split() for x in r)						)
		_UpperCamelCase :		List[str]								=    ' '.join(UpperCAmelCase_						).split()
		args_exp.extend(['--bs', str(args.bs						)]						)  # in case we need to reduce its size due to CUDA OOM
		_UpperCamelCase :		Union[str, Any]								=    args_normal + args_exp

		# XXX: need to trap CUDA OOM and lower args.bs if that happens and retry

		_UpperCamelCase :		Any								=    run_generate(verbose=UpperCAmelCase_						)
		# make sure scores are first in the table
		_UpperCamelCase :		Optional[int]								=    OrderedDict()
		for score in task_score_names[task]:
			_UpperCamelCase :		Tuple								=    scores[score]
		result.update(UpperCAmelCase_						)
		results.append(UpperCAmelCase_						)

		# find widest entries
		for k, v in result.items():
			_UpperCamelCase :		Tuple								=    len(str(UpperCAmelCase_						)						)
			if l > col_widths[k]:
				_UpperCamelCase :		str								=    l

	_UpperCamelCase :		Optional[int]								=    sorted(UpperCAmelCase_				,							key=operator.itemgetter(*task_score_names[task]						)				,							reverse=UpperCAmelCase_						)
	print(' | '.join([F'''{col:{col_widths[col]}}''' for col in col_names]						)						)
	print(' | '.join([F'''{"-"*col_widths[col]}''' for col in col_names]						)						)
	for row in results_sorted:
		print(' | '.join([F'''{row[col]:{col_widths[col]}}''' for col in col_names]						)						)

	_UpperCamelCase :		Optional[int]								=    results_sorted[0]
	for score in task_score_names[task]:
		del best[score]
	_UpperCamelCase :		str								=    [F'''--{k} {v}''' for k, v in best.items()]
	_UpperCamelCase :		Optional[Any]								=    ['--bs', str(args.bs						)]
	if args.info:
		print(F'''\nInfo: {args.info}'''						)
	print('\nBest score args:'						)
	print(' '.join(args_main + best_args + dyn_args						)						)

	return results_sorted


if __name__ == "__main__":
							# Usage:
							# [normal-run_eval_search.py cmd plus] \
							# --search="num_beams=1:5:10 length_penalty=0.8:1:1.2 early_stopping=true:false"
							#
							# Example:
							# PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval_search.py $MODEL_NAME \
							# $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target \
							# --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation \
							# --search="num_beams=1:5:10 length_penalty=0.8:1:1.2 early_stopping=true:false"
							run_search()




