


"""simple docstring"""



import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch
def      UpperCAmelCase     (       _lowercase :   Union[str, Any]	)						-> Union[str, Any]:


	"""simple docstring"""




	lowerCAmelCase_      							=			os.path.join(args.tf_model_dir			,   '''parameters.json'''	)
	lowerCAmelCase_      							=			json.loads(open(_lowercase	).read()	)
	if not params:
		raise ValueError(
		    F"""It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file."""	)
	if not args.output.endswith('''.pt'''	):
		lowerCAmelCase_      							=			args.output + '''.pt'''
	lowerCAmelCase_      							=			OrderedDict()
	with tf.device('''/CPU:0'''	):
		lowerCAmelCase_      							=			tf.train.load_checkpoint(args.tf_model_dir	)
		lowerCAmelCase_      							=			reader.get_variable_to_shape_map()
		for key_name in shapes.keys():
			lowerCAmelCase_      							=			reader.get_tensor(_lowercase	).astype(np.floataa	)
			if key_name.endswith('''/adam_m'''	) or key_name.endswith('''/adam_v'''	):
				continue
			if key_name.startswith('''pasts/'''	):
				if key_name.startswith('''pasts/mlp'''	):
					lowerCAmelCase_      							=			int(key_name[9]	)
				elif key_name.startswith('''pasts/out'''	):
					lowerCAmelCase_      							=			8
				lowerCAmelCase_      							=			'''model.sqout.%d.weight''' % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
				lowerCAmelCase_      							=			vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
				lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name.startswith('''model/moe'''	):
				lowerCAmelCase_      							=			int(key_name[9:].split('''/'''	)[0]	)
				if key_name.endswith('''/switch_gating/kernel'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.mlp.router.classifier.weight''' % player
					lowerCAmelCase_      							=			vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/softmlp/kernel'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.soft_bypass_mlp.weight''' % player
					lowerCAmelCase_      							=			vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/wo/kernel'''	) or key_name.endswith('''/wi/kernel'''	):
					lowerCAmelCase_      							=			key_name[-9:-7]
					for i in range(1_6	):
						lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight''' % (player, i, nlayer)
						lowerCAmelCase_      							=			(
						    vnp[i].transpose([1, 0]	).copy()
						)  # In Mesh-Tensorflow, it is one array, so it is divided
						lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name.startswith('''model/mlp'''	):
				lowerCAmelCase_      							=			int(key_name[9:].split('''/'''	)[0]	)
				if key_name.endswith('''/p1/kernel'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.mlp.wi.weight''' % player
					lowerCAmelCase_      							=			vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/p1/bias'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.mlp.wi.bias''' % player
					lowerCAmelCase_      							=			vnp.copy()  # same because it is one dimensional
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/p2/kernel'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.mlp.wo.weight''' % player
					lowerCAmelCase_      							=			vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/p2/bias'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.mlp.wo.bias''' % player
					lowerCAmelCase_      							=			vnp.copy()  # same because it is one dimensional
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name.startswith('''model/ln'''	):
				lowerCAmelCase_      							=			int(key_name[8:].split('''/'''	)[0]	)
				if key_name.endswith('''/b'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.norm.bias''' % player
					lowerCAmelCase_      							=			vnp.copy()  # same because it is one dimensional
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/g'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.feed_forward.norm.weight''' % player
					lowerCAmelCase_      							=			vnp.copy()  # same because it is one dimensional
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name.startswith('''model/att'''	):
				lowerCAmelCase_      							=			int(key_name[9:].split('''/'''	)[0]	)
				if key_name.endswith('''/qkv/kernel'''	):
					lowerCAmelCase_      							=			vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
					lowerCAmelCase_      							=			state[:, 0, :, :]
					lowerCAmelCase_      							=			state[:, 1, :, :]
					lowerCAmelCase_      							=			state[:, 2, :, :]
					lowerCAmelCase_      							=			(
					    state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]]	)
					    .transpose([1, 0]	)
					    .copy()
					)  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			(
					    state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]]	)
					    .transpose([1, 0]	)
					    .copy()
					)  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			(
					    state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]]	)
					    .transpose([1, 0]	)
					    .copy()
					)  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			'''model.blocks.%d.self_attn.self_attn.q_proj.weight''' % player
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
					lowerCAmelCase_      							=			'''model.blocks.%d.self_attn.self_attn.k_proj.weight''' % player
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
					lowerCAmelCase_      							=			'''model.blocks.%d.self_attn.self_attn.v_proj.weight''' % player
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/o/kernel'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.self_attn.self_attn.out_proj.weight''' % player
					lowerCAmelCase_      							=			(
					    vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]	).transpose([1, 0]	).copy()
					)  # Mesh-Tensorflow is a diagonal matrix
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name.startswith('''model/an'''	):
				lowerCAmelCase_      							=			int(key_name[8:].split('''/'''	)[0]	)
				if key_name.endswith('''/b'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.self_attn.norm.bias''' % player
					lowerCAmelCase_      							=			vnp.copy()  # same because it is one dimensional
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				elif key_name.endswith('''/g'''	):
					lowerCAmelCase_      							=			'''model.blocks.%d.self_attn.norm.weight''' % player
					lowerCAmelCase_      							=			vnp.copy()  # same because it is one dimensional
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif (
			    key_name.startswith('''model/wte'''	)
			    or key_name.startswith('''model/wpe'''	)
			    or key_name.startswith('''model/ete'''	)
			):
				lowerCAmelCase_      							=			{'''wte''': '''embed_tokens''', '''wpe''': '''position_embeddings''', '''ete''': '''extra_position_embeddings'''}[
				    key_name[-3:]
				]
				lowerCAmelCase_      							=			'''model.%s.weight''' % nlayer
				lowerCAmelCase_      							=			vnp.copy()  # same in embedded
				lowerCAmelCase_      							=			torch.tensor(_lowercase	)
				if key_name.startswith('''model/wte'''	):
					lowerCAmelCase_      							=			'''lm_head.weight'''
					lowerCAmelCase_      							=			vnp.copy()  # same in embedded
					lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name.startswith('''model/wob'''	):
				lowerCAmelCase_      							=			'''final_logits_bias'''
				lowerCAmelCase_      							=			vnp.copy()  # same in embedded
				lowerCAmelCase_      							=			state.reshape((1, -1)	)
				lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name == "model/dense/kernel":
				lowerCAmelCase_      							=			'''model.last_project.weight'''
				lowerCAmelCase_      							=			vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
				lowerCAmelCase_      							=			torch.tensor(_lowercase	)
			elif key_name == "model/dense_1/bias":
				lowerCAmelCase_      							=			'''model.last_project.bias'''
				lowerCAmelCase_      							=			vnp.copy()  # same because it is one dimensional
				lowerCAmelCase_      							=			torch.tensor(_lowercase	)
	torch.save(_lowercase			,   args.output	)


if __name__ == "__main__":
		lowercase_							=							argparse.ArgumentParser(
		    description='model converter.', formatter_class=argparse.ArgumentDefaultsHelpFormatter
		)
		parser.add_argument('--tf_model_dir', metavar='PATH', type=str, required=True, help='import model')
		parser.add_argument('--output', metavar='PATH', type=str, required=True, help='output model')
		lowercase_							=							parser.parse_args()
		convert_tf_gptsan_to_pt(args)



"""simple docstring"""



import os
from itertools import chain
from random import randrange, shuffle

import pytest

from .sola import PokerHand

lowercase_							=							(
    '4S 3H 2C 7S 5H',
    '9D 8H 2C 6S 7H',
    '2D 6D 9D TH 7D',
    'TC 8C 2S JH 6C',
    'JH 8S TH AH QH',
    'TS KS 5S 9S AC',
    'KD 6S 9D TH AD',
    'KS 8D 4D 9S 4S',  # pair
    '8C 4S KH JS 4D',  # pair
    'QH 8H KD JH 8S',  # pair
    'KC 4H KS 2H 8D',  # pair
    'KD 4S KC 3H 8S',  # pair
    'AH 8S AS KC JH',  # pair
    '3H 4C 4H 3S 2H',  # 2 pairs
    '5S 5D 2C KH KH',  # 2 pairs
    '3C KH 5D 5S KH',  # 2 pairs
    'AS 3C KH AD KH',  # 2 pairs
    '7C 7S 3S 7H 5S',  # 3 of a kind
    '7C 7S KH 2H 7H',  # 3 of a kind
    'AC KH QH AH AS',  # 3 of a kind
    '2H 4D 3C AS 5S',  # straight (low ace)
    '3C 5C 4C 2C 6H',  # straight
    '6S 8S 7S 5H 9H',  # straight
    'JS QS 9H TS KH',  # straight
    'QC KH TS JS AH',  # straight (high ace)
    '8C 9C 5C 3C TC',  # flush
    '3S 8S 9S 5S KS',  # flush
    '4C 5C 9C 8C KC',  # flush
    'JH 8H AH KH QH',  # flush
    '3D 2H 3H 2C 2D',  # full house
    '2H 2C 3S 3H 3D',  # full house
    'KH KC 3S 3H 3D',  # full house
    'JC 6H JS JD JH',  # 4 of a kind
    'JC 7H JS JD JH',  # 4 of a kind
    'JC KH JS JD JH',  # 4 of a kind
    '2S AS 4S 5S 3S',  # straight flush (low ace)
    '2D 6D 3D 4D 5D',  # straight flush
    '5C 6C 3C 7C 4C',  # straight flush
    'JH 9H TH KH QH',  # straight flush
    'JH AH TH KH QH',  # royal flush (high ace straight flush)
)

lowercase_							=							(
    ('2H 3H 4H 5H 6H', 'KS AS TS QS JS', 'Loss'),
    ('2H 3H 4H 5H 6H', 'AS AD AC AH JD', 'Win'),
    ('AS AH 2H AD AC', 'JS JD JC JH 3D', 'Win'),
    ('2S AH 2H AS AC', 'JS JD JC JH AD', 'Loss'),
    ('2S AH 2H AS AC', '2H 3H 5H 6H 7H', 'Win'),
    ('AS 3S 4S 8S 2S', '2H 3H 5H 6H 7H', 'Win'),
    ('2H 3H 5H 6H 7H', '2S 3H 4H 5S 6C', 'Win'),
    ('2S 3H 4H 5S 6C', '3D 4C 5H 6H 2S', 'Tie'),
    ('2S 3H 4H 5S 6C', 'AH AC 5H 6H AS', 'Win'),
    ('2S 2H 4H 5S 4C', 'AH AC 5H 6H AS', 'Loss'),
    ('2S 2H 4H 5S 4C', 'AH AC 5H 6H 7S', 'Win'),
    ('6S AD 7H 4S AS', 'AH AC 5H 6H 7S', 'Loss'),
    ('2S AH 4H 5S KC', 'AH AC 5H 6H 7S', 'Loss'),
    ('2S 3H 6H 7S 9C', '7H 3C TH 6H 9S', 'Loss'),
    ('4S 5H 6H TS AC', '3S 5H 6H TS AC', 'Win'),
    ('2S AH 4H 5S 6C', 'AD 4C 5H 6H 2C', 'Tie'),
    ('AS AH 3H AD AC', 'AS AH 2H AD AC', 'Win'),
    ('AH AC 5H 5C QS', 'AH AC 5H 5C KS', 'Loss'),
    ('AH AC 5H 5C QS', 'KH KC 5H 5C QS', 'Win'),
    ('7C 7S KH 2H 7H', '3C 3S AH 2H 3H', 'Win'),
    ('3C 3S AH 2H 3H', '7C 7S KH 2H 7H', 'Loss'),
    ('6H 5H 4H 3H 2H', '5H 4H 3H 2H AH', 'Win'),
    ('5H 4H 3H 2H AH', '5H 4H 3H 2H AH', 'Tie'),
    ('5H 4H 3H 2H AH', '6H 5H 4H 3H 2H', 'Loss'),
    ('AH AD KS KC AC', 'AH KD KH AC KC', 'Win'),
    ('2H 4D 3C AS 5S', '2H 4D 3C 6S 5S', 'Loss'),
    ('2H 3S 3C 3H 2S', '3S 3C 2S 2H 2D', 'Win'),
    ('4D 6D 5D 2D JH', '3S 8S 3H TC KH', 'Loss'),
    ('4S 6C 8S 3S 7S', 'AD KS 2D 7D 7C', 'Loss'),
    ('6S 4C 7H 8C 3H', '5H JC AH 9D 9C', 'Loss'),
    ('9D 9H JH TC QH', '3C 2S JS 5C 7H', 'Win'),
    ('2H TC 8S AD 9S', '4H TS 7H 2C 5C', 'Win'),
    ('9D 3S 2C 7S 7C', 'JC TD 3C TC 9H', 'Loss'),
)

lowercase_							=							(
    ('2H 3H 4H 5H 6H', True),
    ('AS AH 2H AD AC', False),
    ('2H 3H 5H 6H 7H', True),
    ('KS AS TS QS JS', True),
    ('8H 9H QS JS TH', False),
    ('AS 3S 4S 8S 2S', True),
)

lowercase_							=							(
    ('2H 3H 4H 5H 6H', True),
    ('AS AH 2H AD AC', False),
    ('2H 3H 5H 6H 7H', False),
    ('KS AS TS QS JS', True),
    ('8H 9H QS JS TH', True),
)

lowercase_							=							(
    ('2H 4D 3C AS 5S', True, [5, 4, 3, 2, 14]),
    ('2H 5D 3C AS 5S', False, [14, 5, 5, 3, 2]),
    ('JH QD KC AS TS', False, [14, 13, 12, 11, 10]),
    ('9D 3S 2C 7S 7C', False, [9, 7, 7, 3, 2]),
)

lowercase_							=							(
    ('JH AH TH KH QH', 0),
    ('JH 9H TH KH QH', 0),
    ('JC KH JS JD JH', 7),
    ('KH KC 3S 3H 3D', 6),
    ('8C 9C 5C 3C TC', 0),
    ('JS QS 9H TS KH', 0),
    ('7C 7S KH 2H 7H', 3),
    ('3C KH 5D 5S KH', 2),
    ('QH 8H KD JH 8S', 1),
    ('2D 6D 9D TH 7D', 0),
)

lowercase_							=							(
    ('JH AH TH KH QH', 23),
    ('JH 9H TH KH QH', 22),
    ('JC KH JS JD JH', 21),
    ('KH KC 3S 3H 3D', 20),
    ('8C 9C 5C 3C TC', 19),
    ('JS QS 9H TS KH', 18),
    ('7C 7S KH 2H 7H', 17),
    ('3C KH 5D 5S KH', 16),
    ('QH 8H KD JH 8S', 15),
    ('2D 6D 9D TH 7D', 14),
)
def      UpperCAmelCase     (       )						-> List[str]:


	"""simple docstring"""




	lowerCAmelCase_							,	lowerCAmelCase_      							=			randrange(len(_lowercase	)	), randrange(len(_lowercase	)	)
	lowerCAmelCase_      							=			['''Loss''', '''Tie''', '''Win'''][(play >= oppo) + (play > oppo)]
	lowerCAmelCase_							,	lowerCAmelCase_      							=			SORTED_HANDS[play], SORTED_HANDS[oppo]
	return hand, other, expected
def      UpperCAmelCase     (       _lowercase :   int = 1_0_0	)						-> int:


	"""simple docstring"""




	return (generate_random_hand() for _ in range(_lowercase	))
@pytest.mark.parametrize('''hand, expected'''			,   _lowercase	)
def      UpperCAmelCase     (       _lowercase :   Union[str, Any]			,   _lowercase :   Optional[int]	)						-> str:


	"""simple docstring"""




	assert PokerHand(_lowercase	)._is_flush() == expected
@pytest.mark.parametrize('''hand, expected'''			,   _lowercase	)
def      UpperCAmelCase     (       _lowercase :   List[Any]			,   _lowercase :   List[str]	)						-> Optional[Any]:


	"""simple docstring"""




	assert PokerHand(_lowercase	)._is_straight() == expected
@pytest.mark.parametrize('''hand, expected, card_values'''			,   _lowercase	)
def      UpperCAmelCase     (       _lowercase :   str			,   _lowercase :   Any			,   _lowercase :   Union[str, Any]	)						-> str:


	"""simple docstring"""




	lowerCAmelCase_      							=			PokerHand(_lowercase	)
	assert player._is_five_high_straight() == expected
	assert player._card_values == card_values
@pytest.mark.parametrize('''hand, expected'''			,   _lowercase	)
def      UpperCAmelCase     (       _lowercase :   Any			,   _lowercase :   Union[str, Any]	)						-> Optional[Any]:


	"""simple docstring"""




	assert PokerHand(_lowercase	)._is_same_kind() == expected
@pytest.mark.parametrize('''hand, expected'''			,   _lowercase	)
def      UpperCAmelCase     (       _lowercase :   Tuple			,   _lowercase :   int	)						-> str:


	"""simple docstring"""




	assert PokerHand(_lowercase	)._hand_type == expected
@pytest.mark.parametrize('''hand, other, expected'''			,   _lowercase	)
def      UpperCAmelCase     (       _lowercase :   int			,   _lowercase :   str			,   _lowercase :   Tuple	)						-> Union[str, Any]:


	"""simple docstring"""




	assert PokerHand(_lowercase	).compare_with(PokerHand(_lowercase	)	) == expected
@pytest.mark.parametrize('''hand, other, expected'''			,   generate_random_hands()	)
def      UpperCAmelCase     (       _lowercase :   Any			,   _lowercase :   Dict			,   _lowercase :   Dict	)						-> Union[str, Any]:


	"""simple docstring"""




	assert PokerHand(_lowercase	).compare_with(PokerHand(_lowercase	)	) == expected
def      UpperCAmelCase     (       )						-> List[str]:


	"""simple docstring"""




	lowerCAmelCase_      							=			[PokerHand(_lowercase	) for hand in SORTED_HANDS]
	lowerCAmelCase_      							=			poker_hands.copy()
	shuffle(_lowercase	)
	lowerCAmelCase_      							=			chain(sorted(_lowercase	)	)
	for index, hand in enumerate(_lowercase	):
		assert hand == poker_hands[index]
def      UpperCAmelCase     (       )						-> Any:


	"""simple docstring"""




	lowerCAmelCase_      							=			[PokerHand('''2D AC 3H 4H 5S'''	), PokerHand('''2S 3H 4H 5S 6C'''	)]
	pokerhands.sort(reverse=_lowercase	)
	assert pokerhands[0].__str__() == "2S 3H 4H 5S 6C"
def      UpperCAmelCase     (       )						-> List[str]:


	"""simple docstring"""




	lowerCAmelCase_      							=			PokerHand('''2C 4S AS 3D 5C'''	)
	lowerCAmelCase_      							=			True
	lowerCAmelCase_      							=			[5, 4, 3, 2, 1_4]
	for _ in range(1_0	):
		assert pokerhand._is_five_high_straight() == expected
		assert pokerhand._card_values == expected_card_values





def      UpperCAmelCase     (       )						-> Dict:


	"""simple docstring"""




	lowerCAmelCase_      							=			0
	lowerCAmelCase_      							=			os.path.abspath(os.path.dirname(_lowercase	)	)
	lowerCAmelCase_      							=			os.path.join(_lowercase			,   '''poker_hands.txt'''	)
	with open(_lowercase	) as file_hand:
		for line in file_hand:
			lowerCAmelCase_      							=			line[:1_4].strip()
			lowerCAmelCase_      							=			line[1_5:].strip()
			lowerCAmelCase_							,	lowerCAmelCase_      							=			PokerHand(_lowercase	), PokerHand(_lowercase	)
			lowerCAmelCase_      							=			player.compare_with(_lowercase	)
			if output == "Win":
				answer += 1
	assert answer == 3_7_6