
'''simple docstring'''






import collections
import json
import math
import os
import re
import time
from fnmatch import fnmatch
from typing import Dict

import requests
from slack_sdk import WebClient


A			      =      WebClient(token=os.environ['CI_SLACK_BOT_TOKEN'])





def    UpperCAmelCase							(  UpperCAmelCase__							:							Tuple):
				lowerCamelCase					:			List[str]									= test_results.split(' ')

				lowerCamelCase					:			Optional[Any]									= 0
				lowerCamelCase					:			Any									= 0

				# When the output is short enough, the output is surrounded by = signs: "== OUTPUT =="
				# When it is too long, those signs are not present.
				lowerCamelCase					:			Optional[Any]									= expressions[-2] if '=' in expressions[-1] else expressions[-1]

				for i, expression in enumerate(UpperCAmelCase__):
								if "failed" in expression:
												failed += int(expressions[i - 1])
								if "passed" in expression:
												success += int(expressions[i - 1])

				return failed, success, time_spent





def    UpperCAmelCase							(  UpperCAmelCase__							:							Optional[int]):
				lowerCamelCase					:			Union[str, Any]									= {}
				lowerCamelCase					:			str									= None
				lowerCamelCase					:			Dict									= False
				for line in failures_short_lines.split('\n'):
								if re.search(R'_ \[doctest\]'       ,							UpperCAmelCase__):
												lowerCamelCase					:			Any									= True
												lowerCamelCase					:			Dict									= line.split(' ')[2]
								elif in_error and not line.split(' ')[0].isdigit():
												lowerCamelCase					:			int									= line
												lowerCamelCase					:			Dict									= False

				return failures




class     __snake_case		:
					def __init__(     self,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= title

									lowerCamelCase					:			Optional[int]									= doc_test_results['time_spent'].split(','      )[0]
									lowerCamelCase					:			Dict									= doc_test_results['success']
									lowerCamelCase					:			str									= doc_test_results['failures']
									lowerCamelCase					:			str									= self.n_success + self.n_failures

									# Failures and success of the modeling tests
									lowerCamelCase					:			str									= doc_test_results
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= [self._time_spent]
									lowerCamelCase					:			List[str]									= 0

									for time in time_spent:
													lowerCamelCase					:			str									= time.split(':'      )

													# Time can be formatted as xx:xx:xx, as .xx, or as x.xx if the time spent was less than a minute.
													if len(A      ) == 1:
																	lowerCamelCase					:			int									= [0, 0, time_parts[0]]

													lowerCamelCase			, lowerCamelCase			, lowerCamelCase					:			Optional[int]									= int(time_parts[0]      ), int(time_parts[1]      ), float(time_parts[2]      )
													total_secs += hours * 3600 + minutes * 60 + seconds

									lowerCamelCase			, lowerCamelCase			, lowerCamelCase					:			Any									= total_secs // 3600, (total_secs % 3600) // 60, total_secs % 60
									return F'''{int(A      )}h{int(A      )}m{int(A      )}s'''
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return {"type": "header", "text": {"type": "plain_text", "text": self.title}}
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return {
									    "type": "section",
									    "text": {
									        "type": "plain_text",
									        "text": F'''🌞 There were no failures: all {self.n_tests} tests passed. The suite ran in {self.time}.''',
									        "emoji": True,
									    },
									    "accessory": {
									        "type": "button",
									        "text": {"type": "plain_text", "text": "Check Action results", "emoji": True},
									        "url": F'''https://github.com/huggingface/transformers/actions/runs/{os.environ['GITHUB_RUN_ID']}''',
									    },
									}
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									return {
									    "type": "section",
									    "text": {
									        "type": "plain_text",
									        "text": (
									            F'''There were {self.n_failures} failures, out of {self.n_tests} tests.\nThe suite ran in'''
									            F''' {self.time}.'''
									        ),
									        "emoji": True,
									    },
									    "accessory": {
									        "type": "button",
									        "text": {"type": "plain_text", "text": "Check Action results", "emoji": True},
									        "url": F'''https://github.com/huggingface/transformers/actions/runs/{os.environ['GITHUB_RUN_ID']}''',
									    },
									}
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= 40
									lowerCamelCase					:			int									= {k: v['failed'] for k, v in doc_test_results.items() if isinstance(A,    A      )}

									lowerCamelCase					:			Dict									= ''
									for category, failures in category_failures.items():
													if len(A      ) == 0:
																	continue

													if report != "":
																	report += "\n\n"

													report += F'''*{category} failures*:'''.ljust(line_length // 2      ).rjust(line_length // 2      ) + "\n"
													report += "`"
													report += "`\n`".join(A      )
													report += "`"

									return {
									    "type": "section",
									    "text": {
									        "type": "mrkdwn",
									        "text": F'''The following examples had failures:\n\n\n{report}\n''',
									    },
									}
					@property
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Dict									= [self.header]

									if self.n_failures > 0:
													blocks.append(self.failures      )

									if self.n_failures > 0:
													blocks.extend([self.category_failures]      )

									if self.n_failures == 0:
													blocks.append(self.no_failures      )

									return json.dumps(A      )
					@staticmethod
					def     UpperCAmelCase_							(     ):

									"""simple docstring"""


									lowerCamelCase					:			str									= [
									    {
									        'type': 'section',
									        'text': {
									            'type': 'plain_text',
									            'text': 'There was an issue running the tests.',
									        },
									        'accessory': {
									            'type': 'button',
									            'text': {'type': 'plain_text', 'text': 'Check Action results', 'emoji': True},
									            'url': F'''https://github.com/huggingface/transformers/actions/runs/{os.environ['GITHUB_RUN_ID']}''',
									        },
									    }
									]

									print('Sending the following payload'      )
									print(json.dumps({'blocks': json.loads(A      )}      )      )

									client.chat_postMessage(
									    channel=os.environ['CI_SLACK_CHANNEL_ID_DAILY'],    text='There was an issue running the tests.',    blocks=A,    )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									print('Sending the following payload'      )
									print(json.dumps({'blocks': json.loads(self.payload      )}      )      )

									lowerCamelCase					:			Optional[Any]									= F'''{self.n_failures} failures out of {self.n_tests} tests,''' if self.n_failures else 'All tests passed.'

									lowerCamelCase					:			List[Any]									= client.chat_postMessage(
									    channel=os.environ['CI_SLACK_CHANNEL_ID_DAILY'],    blocks=self.payload,    text=A,    )
					def     UpperCAmelCase_							(     self,    A,    A,    A,    A      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= ''
									for key, value in failures.items():
													lowerCamelCase					:			Dict									= value[:200] + ' [Truncated]' if len(A      ) > 250 else value
													failures_text += F'''*{key}*\n_{value}_\n\n'''

									lowerCamelCase					:			Tuple									= job_name
									lowerCamelCase					:			List[Any]									= {'type': 'section', 'text': {'type': 'mrkdwn', 'text': text}}

									if job_link is not None:
													lowerCamelCase					:			Any									= {
													    'type': 'button',
													    'text': {'type': 'plain_text', 'text': 'GitHub Action job', 'emoji': True},
													    'url': job_link,
													}

									return [
									    {"type": "header", "text": {"type": "plain_text", "text": title.upper(), "emoji": True}},
									    content,
									    {"type": "section", "text": {"type": "mrkdwn", "text": failures_text}},
									]




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									if self.thread_ts is None:
													raise ValueError('Can only post reply if a post has been made.'      )

									lowerCamelCase					:			Tuple									= self.doc_test_results.pop('job_link'      )
									self.doc_test_results.pop('failures'      )
									self.doc_test_results.pop('success'      )
									self.doc_test_results.pop('time_spent'      )

									lowerCamelCase					:			Dict									= sorted(self.doc_test_results.items(),    key=lambda A      : t[0]      )
									for job, job_result in sorted_dict:
													if len(job_result['failures']      ):
																	lowerCamelCase					:			str									= F'''*Num failures* :{len(job_result['failed']      )} \n'''
																	lowerCamelCase					:			Optional[int]									= job_result['failures']
																	lowerCamelCase					:			Optional[int]									= self.get_reply_blocks(A,    A,    A,    text=A      )

																	print('Sending the following reply'      )
																	print(json.dumps({'blocks': blocks}      )      )

																	client.chat_postMessage(
																	    channel=os.environ['CI_SLACK_CHANNEL_ID_DAILY'],    text=F'''Results for {job}''',    blocks=A,    thread_ts=self.thread_ts['ts'],    )

																	time.sleep(1      )





def    UpperCAmelCase							(  ):
				lowerCamelCase					:			List[str]									= os.environ['GITHUB_RUN_ID']
				lowerCamelCase					:			Tuple									= F'''https://api.github.com/repos/huggingface/transformers/actions/runs/{run_id}/jobs?per_page=100'''
				lowerCamelCase					:			Optional[int]									= requests.get(UpperCAmelCase__).json()
				lowerCamelCase					:			Dict									= {}

				try:
								jobs.update({job['name']: job['html_url'] for job in result['jobs']})
								lowerCamelCase					:			Optional[int]									= math.ceil((result['total_count'] - 1_00) / 1_00)

								for i in range(UpperCAmelCase__):
												lowerCamelCase					:			List[Any]									= requests.get(url + F'''&page={i + 2}''').json()
												jobs.update({job['name']: job['html_url'] for job in result['jobs']})

								return jobs
				except Exception as e:
								print('Unknown error, could not fetch links.'       ,							UpperCAmelCase__)

				return {}





def    UpperCAmelCase							(  UpperCAmelCase__							:							str):
				lowerCamelCase					:			Any									= {}

				if os.path.exists(UpperCAmelCase__):
								lowerCamelCase					:			List[str]									= os.listdir(UpperCAmelCase__)
								for file in files:
												try:
																with open(os.path.join(UpperCAmelCase__       ,							UpperCAmelCase__)       ,							encoding='utf-8') as f:
																				lowerCamelCase					:			Tuple									= f.read()
												except UnicodeDecodeError as e:
																raise ValueError(F'''Could not open {os.path.join(UpperCAmelCase__       ,							UpperCAmelCase__)}.''') from e

				return _artifact





def    UpperCAmelCase							(  ):




				class     __snake_case		:
									def __init__(     self,    A      ):

													"""simple docstring"""


													lowerCamelCase					:			Union[str, Any]									= name
													lowerCamelCase					:			Optional[Any]									= []
									def __str__(     self      ):

													"""simple docstring"""


													return self.name




									def     UpperCAmelCase_							(     self,    A      ):

													"""simple docstring"""


													self.paths.append({'name': self.name, 'path': path}      )





				lowerCamelCase					:			Dict[str, Artifact]									= {}

				lowerCamelCase					:			Optional[Any]									= filter(os.path.isdir       ,							os.listdir())
				for directory in directories:
								lowerCamelCase					:			Optional[Any]									= directory
								if artifact_name not in _available_artifacts:
												lowerCamelCase					:			Any									= Artifact(UpperCAmelCase__)

												_available_artifacts[artifact_name].add_path(UpperCAmelCase__)

				return _available_artifacts


if __name__ == "__main__":
					A			      =      get_job_links()
					A			      =      retrieve_available_artifacts()

					A			      =      collections.OrderedDict(
					    [
					        ('*.py', 'API Examples'),
					        ('*.md', 'MD Examples'),
					    ]
					)

					# This dict will contain all the information relative to each doc test category:
					# - failed: list of failed tests
					# - failures: dict in the format 'test': 'error_message'
					A			      =      {
					    v: {
					        'failed': [],
					        'failures': {},
					    }
					    for v in docs.values()
					}

					# Link to the GitHub Action job
					A			      =      github_actions_job_links.get('run_doctests')

					A			      =      available_artifacts['doc_tests_gpu_test_reports'].paths[0]
					A			      =      retrieve_artifact(artifact_path['name'])
					if "stats" in artifact:
										A      ,			A      ,			A			      =      handle_test_results(artifact['stats'])
										A			      =      failed
										A			      =      success
										A			      =      time_spent[1:-1] + ', '

										A			      =      extract_first_line_failure(artifact['failures_short'])
										for line in artifact["summary_short"].split('\n'):
															if re.search('FAILED', line):
																				A			      =      line.replace('FAILED ', '')
																				A			      =      line.split()[0].replace('\n', '')

																				if "::" in line:
																									A      ,			A			      =      line.split('::')
																				else:
																									A      ,			A			      =      line, line

																				for file_regex in docs.keys():
																									if fnmatch(file_path, file_regex):
																														A			      =      docs[file_regex]
																														doc_test_results[category]["failed"].append(test)

																														A			      =      all_failures[test] if test in all_failures else 'N/A'
																														A			      =      failure
																														break

					A			      =      Message('🤗 Results of the doc tests.', doc_test_results)
					message.post()
					message.post_reply()


'''simple docstring'''






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


A			      =      {
    '169M': 12,
    '430M': 24,
    '1B5': 24,
    '3B': 32,
    '7B': 32,
    '14B': 40,
}

A			      =      {
    '169M': 768,
    '430M': 1024,
    '1B5': 2048,
    '3B': 2560,
    '7B': 4096,
    '14B': 5120,
}





def    UpperCAmelCase							(  UpperCAmelCase__							:							Any):
				lowerCamelCase					:			Any									= list(state_dict.keys())
				for name in state_dict_keys:
								lowerCamelCase					:			List[str]									= state_dict.pop(UpperCAmelCase__)
								# emb -> embedding
								if name.startswith('emb.'):
												lowerCamelCase					:			Dict									= name.replace('emb.'       ,							'embeddings.')
								# ln_0 -> pre_ln (only present at block 0)
								if name.startswith('blocks.0.ln0'):
												lowerCamelCase					:			str									= name.replace('blocks.0.ln0'       ,							'blocks.0.pre_ln')
								# att -> attention
								lowerCamelCase					:			str									= re.sub(R'blocks\.(\d+)\.att'       ,							R'blocks.\1.attention'       ,							UpperCAmelCase__)
								# ffn -> feed_forward
								lowerCamelCase					:			List[Any]									= re.sub(R'blocks\.(\d+)\.ffn'       ,							R'blocks.\1.feed_forward'       ,							UpperCAmelCase__)
								# time_mix_k -> time_mix_key and reshape
								if name.endswith('.time_mix_k'):
												lowerCamelCase					:			Any									= name.replace('.time_mix_k'       ,							'.time_mix_key')
								# time_mix_v -> time_mix_value and reshape
								if name.endswith('.time_mix_v'):
												lowerCamelCase					:			str									= name.replace('.time_mix_v'       ,							'.time_mix_value')
								# time_mix_r -> time_mix_key and reshape
								if name.endswith('.time_mix_r'):
												lowerCamelCase					:			List[Any]									= name.replace('.time_mix_r'       ,							'.time_mix_receptance')

								if name != "head.weight":
												lowerCamelCase					:			Any									= 'rwkv.' + name

								lowerCamelCase					:			Any									= weight
				return state_dict





def    UpperCAmelCase							(  UpperCAmelCase__							:							str       ,							UpperCAmelCase__							:							List[str]       ,							UpperCAmelCase__							:							Optional[int]       ,							UpperCAmelCase__							:							Optional[Any]=None       ,							UpperCAmelCase__							:							Tuple=None       ,							UpperCAmelCase__							:							List[Any]=False       ,							UpperCAmelCase__							:							Dict=None):
				# 1. If possible, build the tokenizer.
				if tokenizer_file is None:
								print('No `--tokenizer_file` provided, we will use the default tokenizer.')
								lowerCamelCase					:			Dict									= 5_02_77
								lowerCamelCase					:			Union[str, Any]									= AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')
				else:
								lowerCamelCase					:			int									= PreTrainedTokenizerFast(tokenizer_file=UpperCAmelCase__)
								lowerCamelCase					:			List[Any]									= len(UpperCAmelCase__)
				tokenizer.save_pretrained(UpperCAmelCase__)

				# 2. Build the config
				lowerCamelCase					:			Tuple									= 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:
																lowerCamelCase					:			Union[str, Any]									= 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}.''')

				lowerCamelCase					:			List[Any]									= RwkvConfig(
				    vocab_size=UpperCAmelCase__       ,							num_hidden_layers=NUM_HIDDEN_LAYERS_MAPPING[size]       ,							hidden_size=HIDEN_SIZE_MAPPING[size]       ,							)
				config.save_pretrained(UpperCAmelCase__)

				# 3. Download model file then convert state_dict
				lowerCamelCase					:			Tuple									= hf_hub_download(UpperCAmelCase__       ,							UpperCAmelCase__)
				lowerCamelCase					:			Any									= torch.load(UpperCAmelCase__       ,							map_location='cpu')
				lowerCamelCase					:			str									= convert_state_dict(UpperCAmelCase__)

				# 4. Split in shards and save
				lowerCamelCase			, lowerCamelCase					:			Optional[int]									= shard_checkpoint(UpperCAmelCase__)
				for shard_file, shard in shards.items():
								torch.save(UpperCAmelCase__       ,							os.path.join(UpperCAmelCase__       ,							UpperCAmelCase__))

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

								# 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.')
								lowerCamelCase					:			List[str]									= list(shards.keys())

								del state_dict
								del shards
								gc.collect()

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

				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.')
								lowerCamelCase					:			List[Any]									= AutoModelForCausalLM.from_pretrained(UpperCAmelCase__)
								model.push_to_hub(UpperCAmelCase__       ,							max_shard_size='2GB')
								tokenizer.push_to_hub(UpperCAmelCase__)


if __name__ == "__main__":
					A			      =      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.',
					)

					A			      =      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,
					)
