



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

import pytest

from .sola import PokerHand

UpperCAmelCase_							:      int      =  (
    "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)
)

UpperCAmelCase_							:      Any      =  (
    ("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"),
)

UpperCAmelCase_							:      Dict      =  (
    ("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),
)

UpperCAmelCase_							:      List[str]      =  (
    ("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),
)

UpperCAmelCase_							:      str      =  (
    ("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]),
)

UpperCAmelCase_							:      List[str]      =  (
    ("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),
)

UpperCAmelCase_							:      Dict      =  (
    ("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						(       )->       Any:





						"""simple docstring"""





						A__						,       A__							=       randrange(len(_A    )    ), randrange(len(_A    )    )
						A__							=       ["Loss", "Tie", "Win"][(play >= oppo) + (play > oppo)]
						A__						,       A__							=       SORTED_HANDS[play], SORTED_HANDS[oppo]
						return hand, other, expected
def  UpperCamelCase						(       _A			:				int = 100    )->       str:





						"""simple docstring"""





						return (generate_random_hand() for _ in range(_A    ))
@pytest.mark.parametrize("hand, expected"			,							_A    )
def  UpperCamelCase						(       _A			:				List[str]			,							_A			:				Any    )->       int:





						"""simple docstring"""





						assert PokerHand(_A    )._is_flush() == expected
@pytest.mark.parametrize("hand, expected"			,							_A    )
def  UpperCamelCase						(       _A			:				int			,							_A			:				str    )->       Dict:





						"""simple docstring"""





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





						"""simple docstring"""





						A__							=       PokerHand(_A    )
						assert player._is_five_high_straight() == expected
						assert player._card_values == card_values
@pytest.mark.parametrize("hand, expected"			,							_A    )
def  UpperCamelCase						(       _A			:				str			,							_A			:				Any    )->       int:





						"""simple docstring"""





						assert PokerHand(_A    )._is_same_kind() == expected
@pytest.mark.parametrize("hand, expected"			,							_A    )
def  UpperCamelCase						(       _A			:				Dict			,							_A			:				List[str]    )->       str:





						"""simple docstring"""





						assert PokerHand(_A    )._hand_type == expected
@pytest.mark.parametrize("hand, other, expected"			,							_A    )
def  UpperCamelCase						(       _A			:				Any			,							_A			:				Any			,							_A			:				Dict    )->       int:





						"""simple docstring"""





						assert PokerHand(_A    ).compare_with(PokerHand(_A    )    ) == expected
@pytest.mark.parametrize("hand, other, expected"			,							generate_random_hands()    )
def  UpperCamelCase						(       _A			:				Optional[Any]			,							_A			:				Optional[int]			,							_A			:				Any    )->       Tuple:





						"""simple docstring"""





						assert PokerHand(_A    ).compare_with(PokerHand(_A    )    ) == expected
def  UpperCamelCase						(       )->       str:





						"""simple docstring"""





						A__							=       [PokerHand(_A    ) for hand in SORTED_HANDS]
						A__							=       poker_hands.copy()
						shuffle(_A    )
						A__							=       chain(sorted(_A    )    )
						for index, hand in enumerate(_A    ):
												assert hand == poker_hands[index]
def  UpperCamelCase						(       )->       Dict:





						"""simple docstring"""





						A__							=       [PokerHand("2D AC 3H 4H 5S"    ), PokerHand("2S 3H 4H 5S 6C"    )]
						pokerhands.sort(reverse=_A    )
						assert pokerhands[0].__str__() == "2S 3H 4H 5S 6C"
def  UpperCamelCase						(       )->       Optional[Any]:





						"""simple docstring"""





						A__							=       PokerHand("2C 4S AS 3D 5C"    )
						A__							=       True
						A__							=       [5, 4, 3, 2, 14]
						for _ in range(10    ):
												assert pokerhand._is_five_high_straight() == expected
												assert pokerhand._card_values == expected_card_values
def  UpperCamelCase						(       )->       Union[str, Any]:





						"""simple docstring"""





						A__							=       0
						A__							=       os.path.abspath(os.path.dirname(_A    )    )
						A__							=       os.path.join(_A			,							"poker_hands.txt"    )
						with open(_A    ) as file_hand:
												for line in file_hand:
																		A__							=       line[:14].strip()
																		A__							=       line[15:].strip()
																		A__						,       A__							=       PokerHand(_A    ), PokerHand(_A    )
																		A__							=       player.compare_with(_A    )
																		if output == "Win":
																								answer += 1
						assert answer == 376










import argparse
import gc
import json
import os
import re

import torch
from huggingface_hub import hf_hub_download

from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizerFast, RwkvConfig
from transformers.modeling_utils import WEIGHTS_INDEX_NAME, shard_checkpoint


UpperCAmelCase_							:      Dict      =  {
    "169M": 12,
    "430M": 24,
    "1B5": 24,
    "3B": 32,
    "7B": 32,
    "14B": 40,
}

UpperCAmelCase_							:      Any      =  {
    "169M": 768,
    "430M": 1_024,
    "1B5": 2_048,
    "3B": 2_560,
    "7B": 4_096,
    "14B": 5_120,
}
def  UpperCamelCase						(       _A			:				Dict    )->       Optional[int]:





						"""simple docstring"""





						A__							=       list(state_dict.keys()    )
						for name in state_dict_keys:
												A__							=       state_dict.pop(_A    )
												# emb -> embedding
												if name.startswith("emb."    ):
																		A__							=       name.replace("emb."			,							"embeddings."    )
												# ln_0 -> pre_ln (only present at block 0)
												if name.startswith("blocks.0.ln0"    ):
																		A__							=       name.replace("blocks.0.ln0"			,							"blocks.0.pre_ln"    )
												# att -> attention
												A__							=       re.sub(R"blocks\.(\d+)\.att"			,							R"blocks.\1.attention"			,							_A    )
												# ffn -> feed_forward
												A__							=       re.sub(R"blocks\.(\d+)\.ffn"			,							R"blocks.\1.feed_forward"			,							_A    )
												# time_mix_k -> time_mix_key and reshape
												if name.endswith(".time_mix_k"    ):
																		A__							=       name.replace(".time_mix_k"			,							".time_mix_key"    )
												# time_mix_v -> time_mix_value and reshape
												if name.endswith(".time_mix_v"    ):
																		A__							=       name.replace(".time_mix_v"			,							".time_mix_value"    )
												# time_mix_r -> time_mix_key and reshape
												if name.endswith(".time_mix_r"    ):
																		A__							=       name.replace(".time_mix_r"			,							".time_mix_receptance"    )

												if name != "head.weight":
																		A__							=       "rwkv." + name

												A__							=       weight
						return state_dict
def  UpperCamelCase						(       _A			:				str			,							_A			:				List[Any]			,							_A			:				List[Any]			,							_A			:				int=None			,							_A			:				List[str]=None			,							_A			:				Dict=False			,							_A			:				List[Any]=None    )->       str:





						"""simple docstring"""





						if tokenizer_file is None:
												print("No `--tokenizer_file` provided, we will use the default tokenizer."    )
												A__							=       50277
												A__							=       AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b"    )
						else:
												A__							=       PreTrainedTokenizerFast(tokenizer_file=_A    )
												A__							=       len(_A    )
						tokenizer.save_pretrained(_A    )

						# 2. Build the config
						A__							=       list(NUM_HIDDEN_LAYERS_MAPPING.keys()    )
						if size is None:
												# Try to infer size from the checkpoint name
												for candidate in possible_sizes:
																		if candidate in checkpoint_file:
																								A__							=       candidate
																								break
												if size is None:
																		raise ValueError("Could not infer the size, please provide it with the `--size` argument."    )
						if size not in possible_sizes:
												raise ValueError(f"""`size` should be one of {possible_sizes}, got {size}."""    )

						A__							=       RwkvConfig(
						    vocab_size=_A			,							num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]			,							hidden_size=HIDEN_SIZE_MAPPING[size]			,							)
						config.save_pretrained(_A    )

						# 3. Download model file then convert state_dict
						A__							=       hf_hub_download(_A			,							_A    )
						A__							=       torch.load(_A			,							map_location="cpu"    )
						A__							=       convert_state_dict(_A    )

						# 4. Split in shards and save
						A__						,       A__							=       shard_checkpoint(_A    )
						for shard_file, shard in shards.items():
												torch.save(_A			,							os.path.join(_A			,							_A    )    )

						if index is not None:
												A__							=       os.path.join(_A			,							_A    )
												# Save the index as well
												with open(_A			,							"w"			,							encoding="utf-8"    ) as f:
																		A__							=       json.dumps(_A			,							indent=2			,							sort_keys=_A    ) + "\n"
																		f.write(_A    )

												# 5. Clean up shards (for some reason the file PyTorch saves take the same space as the whole state_dict
												print(
												    "Cleaning up shards. This may error with an OOM error, it this is the case don't worry you still have converted the model."    )
												A__							=       list(shards.keys()    )

												del state_dict
												del shards
												gc.collect()

												for shard_file in shard_files:
																		A__							=       torch.load(os.path.join(_A			,							_A    )    )
																		torch.save({k: v.cpu().clone() for k, v in state_dict.items()}			,							os.path.join(_A			,							_A    )    )

						del state_dict
						gc.collect()

						if push_to_hub:
												if model_name is None:
																		raise ValueError("Please provide a `model_name` to push the model to the Hub."    )
												A__							=       AutoModelForCausalLM.from_pretrained(_A    )
												model.push_to_hub(_A			,							max_shard_size="2GB"    )
												tokenizer.push_to_hub(_A    )


if __name__ == "__main__":
						UpperCAmelCase_							:      Union[str, Any]      =  argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    "--repo_id", default=None, type=str, required=True, help="Repo ID from which to pull the checkpoint."
						)
						parser.add_argument(
						    "--checkpoint_file", default=None, type=str, required=True, help="Name of the checkpoint file in the repo."
						)
						parser.add_argument(
						    "--output_dir", default=None, type=str, required=True, help="Where to save the converted model."
						)
						parser.add_argument(
						    "--tokenizer_file",
						    default=None,
						    type=str,
						    help="Path to the tokenizer file to use (if not provided, only the model is converted).",
						)
						parser.add_argument(
						    "--size",
						    default=None,
						    type=str,
						    help="Size of the model. Will be inferred from the `checkpoint_file` if not passed.",
						)
						parser.add_argument(
						    "--push_to_hub",
						    action="store_true",
						    help="Push to the Hub the converted model.",
						)
						parser.add_argument(
						    "--model_name",
						    default=None,
						    type=str,
						    help="Name of the pushed model on the Hub, including the username / organization.",
						)

						UpperCAmelCase_							:      Optional[int]      =  parser.parse_args()
						convert_rmkv_checkpoint_to_hf_format(
						    args.repo_id,
						    args.checkpoint_file,
						    args.output_dir,
						    size=args.size,
						    tokenizer_file=args.tokenizer_file,
						    push_to_hub=args.push_to_hub,
						    model_name=args.model_name,
						)





