



'''simple docstring'''

import json
import os
import shutil
import tempfile
import unittest
from multiprocessing import get_context
from pathlib import Path

import datasets
import numpy as np
from datasets import load_dataset
from parameterized import parameterized

from transformers import AutoProcessor
from transformers.models.wavaveca import WavaVecaCTCTokenizer, WavaVecaFeatureExtractor
from transformers.models.wavaveca.tokenization_wavaveca import VOCAB_FILES_NAMES
from transformers.testing_utils import require_pyctcdecode, require_torch, require_torchaudio, slow
from transformers.utils import FEATURE_EXTRACTOR_NAME, is_pyctcdecode_available, is_torch_available

from ..wavaveca.test_feature_extraction_wavaveca import floats_list


if is_pyctcdecode_available():
	from huggingface_hub import snapshot_download
	from pyctcdecode import BeamSearchDecoderCTC

	from transformers.models.wavaveca_with_lm import WavaVecaProcessorWithLM
	from transformers.models.wavaveca_with_lm.processing_wavaveca_with_lm import WavaVecaDecoderWithLMOutput

if is_torch_available():
	from transformers import WavaVecaForCTC






@require_pyctcdecode
class 						a_		(     unittest.TestCase    ):


							def    UpperCamelCase     (				self   :						Tuple			)			->					Dict:
										snake_case:    Optional[int]      							='| <pad> <unk> <s> </s> a b c d e f g h i j k'.split()
										snake_case:    Union[str, Any]      							=dict(zip(a_      ,						range(len(a_			)			)			)			)

										snake_case:    Optional[int]      							={
										    'unk_token': '<unk>',
										    'bos_token': '<s>',
										    'eos_token': '</s>',
										}
										snake_case:    Tuple      							={
										    'feature_size': 1,
										    'padding_value': 0.0,
										    'sampling_rate': 1_6_0_0_0,
										    'return_attention_mask': False,
										    'do_normalize': True,
										}

										snake_case:    int      							=tempfile.mkdtemp()
										snake_case:    List[str]      							=os.path.join(self.tmpdirname      ,						VOCAB_FILES_NAMES['vocab_file']			)
										snake_case:    Tuple      							=os.path.join(self.tmpdirname      ,						a_			)
										with open(self.vocab_file      ,						'w'      ,						encoding='utf-8'			) as fp:
													fp.write(json.dumps(a_			) + '\n'			)

										with open(self.feature_extraction_file      ,						'w'      ,						encoding='utf-8'			) as fp:
													fp.write(json.dumps(a_			) + '\n'			)

										# load decoder from hub
										snake_case:    Dict      							='hf-internal-testing/ngram-beam-search-decoder'


							def    UpperCamelCase     (				self   :						Optional[int]      ,						**a_   :						Union[str, Any]			)			->					Optional[int]:
										snake_case:    List[str]      							=self.add_kwargs_tokens_map.copy()
										kwargs.update(a_			)
										return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname      ,						**a_			)


							def    UpperCamelCase     (				self   :						Union[str, Any]      ,						**a_   :						Any			)			->					Tuple:
										return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname      ,						**a_			)


							def    UpperCamelCase     (				self   :						Optional[Any]      ,						**a_   :						Any			)			->					Optional[Any]:
										return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name      ,						**a_			)


							def    UpperCamelCase     (				self   :						Optional[int]			)			->					List[str]:
										shutil.rmtree(self.tmpdirname			)


							def    UpperCamelCase     (				self   :						Optional[int]			)			->					int:
										snake_case:    List[str]      							=self.get_tokenizer()
										snake_case:    Union[str, Any]      							=self.get_feature_extractor()
										snake_case:    Union[str, Any]      							=self.get_decoder()

										snake_case:    Optional[int]      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										processor.save_pretrained(self.tmpdirname			)
										snake_case:    Tuple      							=WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname			)

										# tokenizer
										self.assertEqual(processor.tokenizer.get_vocab()      ,						tokenizer.get_vocab()			)
										self.assertIsInstance(processor.tokenizer      ,						a_			)

										# feature extractor
										self.assertEqual(processor.feature_extractor.to_json_string()      ,						feature_extractor.to_json_string()			)
										self.assertIsInstance(processor.feature_extractor      ,						a_			)

										# decoder
										self.assertEqual(processor.decoder._alphabet.labels      ,						decoder._alphabet.labels			)
										self.assertEqual(
										    processor.decoder.model_container[decoder._model_key]._unigram_set      ,						decoder.model_container[decoder._model_key]._unigram_set      ,						)
										self.assertIsInstance(processor.decoder      ,						a_			)


							def    UpperCamelCase     (				self   :						List[str]			)			->					str:
										snake_case:    List[Any]      							=WavaVecaProcessorWithLM(
										    tokenizer=self.get_tokenizer()      ,						feature_extractor=self.get_feature_extractor()      ,						decoder=self.get_decoder()			)
										processor.save_pretrained(self.tmpdirname			)

										# make sure that error is thrown when decoder alphabet doesn't match
										snake_case:    Dict      							=WavaVecaProcessorWithLM.from_pretrained(
										    self.tmpdirname      ,						alpha=5.0      ,						beta=3.0      ,						score_boundary=-7.0      ,						unk_score_offset=3			)

										# decoder
										self.assertEqual(processor.language_model.alpha      ,						5.0			)
										self.assertEqual(processor.language_model.beta      ,						3.0			)
										self.assertEqual(processor.language_model.score_boundary      ,						-7.0			)
										self.assertEqual(processor.language_model.unk_score_offset      ,						3			)


							def    UpperCamelCase     (				self   :						Optional[Any]			)			->					List[Any]:
										snake_case:    Optional[Any]      							=self.get_tokenizer()
										# add token to trigger raise
										tokenizer.add_tokens(['xx']			)
										with self.assertRaisesRegex(a_      ,						'include'			):
													WavaVecaProcessorWithLM(
													    tokenizer=a_      ,						feature_extractor=self.get_feature_extractor()      ,						decoder=self.get_decoder()			)


							def    UpperCamelCase     (				self   :						Union[str, Any]			)			->					int:
										snake_case:    Dict      							=self.get_feature_extractor()
										snake_case:    List[str]      							=self.get_tokenizer()
										snake_case:    str      							=self.get_decoder()

										snake_case:    List[Any]      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										snake_case:    Union[str, Any]      							=floats_list((3, 1_0_0_0)			)

										snake_case:    Tuple      							=feature_extractor(a_      ,						return_tensors='np'			)
										snake_case:    Dict      							=processor(a_      ,						return_tensors='np'			)

										for key in input_feat_extract.keys():
													self.assertAlmostEqual(input_feat_extract[key].sum()      ,						input_processor[key].sum()      ,						delta=1E-2			)


							def    UpperCamelCase     (				self   :						List[str]			)			->					int:
										snake_case:    Any      							=self.get_feature_extractor()
										snake_case:    Tuple      							=self.get_tokenizer()
										snake_case:    Dict      							=self.get_decoder()

										snake_case:    str      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										snake_case:    List[Any]      							='This is a test string'

										snake_case:    List[Any]      							=processor(text=a_			)

										snake_case:    Tuple      							=tokenizer(a_			)

										for key in encoded_tok.keys():
													self.assertListEqual(encoded_tok[key]      ,						encoded_processor[key]			)


							def    UpperCamelCase     (				self   :						Optional[Any]      ,						a_   :						Tuple=(2, 1_0, 1_6)      ,						a_   :						Optional[Any]=7_7			)			->					Tuple:
										np.random.seed(a_			)
										return np.random.rand(*a_			)


							def    UpperCamelCase     (				self   :						Tuple			)			->					int:
										snake_case:    Optional[int]      							=self.get_feature_extractor()
										snake_case:    Any      							=self.get_tokenizer()
										snake_case:    List[str]      							=self.get_decoder()

										snake_case:    Union[str, Any]      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										snake_case:    Any      							=self._get_dummy_logits(shape=(1_0, 1_6)      ,						seed=1_3			)

										snake_case:    int      							=processor.decode(a_			)

										snake_case:    List[str]      							=decoder.decode_beams(a_			)[0]

										self.assertEqual(decoded_decoder[0]      ,						decoded_processor.text			)
										self.assertEqual('</s> <s> </s>'      ,						decoded_processor.text			)
										self.assertEqual(decoded_decoder[-2]      ,						decoded_processor.logit_score			)
										self.assertEqual(decoded_decoder[-1]      ,						decoded_processor.lm_score			)


							@parameterized.expand([[None], ['fork'], ['spawn']]			)
							def    UpperCamelCase     (				self   :						int      ,						a_   :						List[str]			)			->					Dict:
										snake_case:    int      							=self.get_feature_extractor()
										snake_case:    Any      							=self.get_tokenizer()
										snake_case:    Tuple      							=self.get_decoder()

										snake_case:    str      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										snake_case:    Dict      							=self._get_dummy_logits()

										# note: pool should be instantiated *after* Wav2Vec2ProcessorWithLM.
										#       otherwise, the LM won't be available to the pool's sub-processes.
										# manual logic used to allow parameterized test for both pool=None and pool=Pool(...)
										if pool_context is None:
													snake_case:    Optional[int]      							=processor.batch_decode(a_			)
										else:
													with get_context(a_			).Pool() as pool:
																snake_case:    int      							=processor.batch_decode(a_      ,						a_			)

										snake_case:    Dict      							=list(a_			)

										with get_context('fork'			).Pool() as p:
													snake_case:    Optional[int]      							=decoder.decode_beams_batch(a_      ,						a_			)

										snake_case							,  snake_case							,  snake_case:    Optional[Any]      							=[], [], []
										for beams in decoded_beams:
													texts_decoder.append(beams[0][0]			)
													logit_scores_decoder.append(beams[0][-2]			)
													lm_scores_decoder.append(beams[0][-1]			)

										self.assertListEqual(a_      ,						decoded_processor.text			)
										self.assertListEqual(['<s> <s> </s>', '<s> <s> <s>']      ,						decoded_processor.text			)
										self.assertListEqual(a_      ,						decoded_processor.logit_score			)
										self.assertListEqual(a_      ,						decoded_processor.lm_score			)


							def    UpperCamelCase     (				self   :						Optional[Any]			)			->					str:
										snake_case:    Any      							=self.get_feature_extractor()
										snake_case:    Tuple      							=self.get_tokenizer()
										snake_case:    int      							=self.get_decoder()

										snake_case:    List[str]      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										snake_case:    List[str]      							=self._get_dummy_logits()

										snake_case:    Optional[Any]      							=1_5
										snake_case:    Dict      							=-2_0.0
										snake_case:    Any      							=-4.0

										snake_case:    List[Any]      							=processor.batch_decode(
										    a_      ,						beam_width=a_      ,						beam_prune_logp=a_      ,						token_min_logp=a_      ,						)
										snake_case:    List[Any]      							=decoded_processor_out.text

										snake_case:    str      							=list(a_			)

										with get_context('fork'			).Pool() as pool:
													snake_case:    Optional[Any]      							=decoder.decode_beams_batch(
													    a_      ,						a_      ,						beam_width=a_      ,						beam_prune_logp=a_      ,						token_min_logp=a_      ,						)

										snake_case:    str      							=[d[0][0] for d in decoded_decoder_out]
										snake_case:    Tuple      							=[d[0][2] for d in decoded_decoder_out]
										snake_case:    Any      							=[d[0][3] for d in decoded_decoder_out]

										self.assertListEqual(a_      ,						a_			)
										self.assertListEqual(['</s> <s> <s>', '<s> <s> <s>']      ,						a_			)

										self.assertTrue(np.array_equal(a_      ,						decoded_processor_out.logit_score			)			)
										self.assertTrue(np.allclose([-2_0.0_5_4, -1_8.4_4_7]      ,						a_      ,						atol=1E-3			)			)

										self.assertTrue(np.array_equal(a_      ,						decoded_processor_out.lm_score			)			)
										self.assertTrue(np.allclose([-1_5.5_5_4, -1_3.9_4_7_4]      ,						a_      ,						atol=1E-3			)			)


							def    UpperCamelCase     (				self   :						List[str]			)			->					List[str]:
										snake_case:    Optional[Any]      							=self.get_feature_extractor()
										snake_case:    List[str]      							=self.get_tokenizer()
										snake_case:    List[Any]      							=self.get_decoder()

										snake_case:    int      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										snake_case:    Any      							=self._get_dummy_logits()

										snake_case:    int      							=2.0
										snake_case:    List[Any]      							=5.0
										snake_case:    Any      							=-2_0.0
										snake_case:    str      							=True

										snake_case:    List[str]      							=processor.batch_decode(
										    a_      ,						alpha=a_      ,						beta=a_      ,						unk_score_offset=a_      ,						lm_score_boundary=a_      ,						)
										snake_case:    Any      							=decoded_processor_out.text

										snake_case:    Union[str, Any]      							=list(a_			)
										decoder.reset_params(
										    alpha=a_      ,						beta=a_      ,						unk_score_offset=a_      ,						lm_score_boundary=a_      ,						)

										with get_context('fork'			).Pool() as pool:
													snake_case:    int      							=decoder.decode_beams_batch(
													    a_      ,						a_      ,						)

										snake_case:    Tuple      							=[d[0][0] for d in decoded_decoder_out]

										self.assertListEqual(a_      ,						a_			)
										self.assertListEqual(['<s> </s> <s> </s> </s>', '</s> </s> <s> </s> </s>']      ,						a_			)
										snake_case:    Any      							=processor.decoder.model_container[processor.decoder._model_key]
										self.assertEqual(lm_model.alpha      ,						2.0			)
										self.assertEqual(lm_model.beta      ,						5.0			)
										self.assertEqual(lm_model.unk_score_offset      ,						-2_0.0			)
										self.assertEqual(lm_model.score_boundary      ,						a_			)


							def    UpperCamelCase     (				self   :						Tuple			)			->					List[Any]:
										snake_case:    Tuple      							=WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)

										snake_case:    int      							=processor.decoder.model_container[processor.decoder._model_key]
										snake_case:    Optional[int]      							=Path(language_model._kenlm_model.path.decode('utf-8'			)			).parent.parent.absolute()

										snake_case:    str      							=os.listdir(a_			)
										snake_case:    Dict      							=['alphabet.json', 'language_model']

										downloaded_decoder_files.sort()
										expected_decoder_files.sort()

										# test that only decoder relevant files from
										# https://huggingface.co/hf-internal-testing/processor_with_lm/tree/main
										# are downloaded and none of the rest (e.g. README.md, ...)
										self.assertListEqual(a_      ,						a_			)


							def    UpperCamelCase     (				self   :						Optional[Any]			)			->					int:
										snake_case:    str      							=snapshot_download('hf-internal-testing/processor_with_lm'			)

										snake_case:    List[str]      							=WavaVecaProcessorWithLM.from_pretrained(a_			)

										snake_case:    Union[str, Any]      							=processor.decoder.model_container[processor.decoder._model_key]
										snake_case:    Dict      							=Path(language_model._kenlm_model.path.decode('utf-8'			)			).parent.parent.absolute()

										snake_case:    int      							=os.listdir(a_			)
										snake_case:    Optional[Any]      							=os.listdir(a_			)

										local_decoder_files.sort()
										expected_decoder_files.sort()

										# test that both decoder form hub and local files in cache are the same
										self.assertListEqual(a_      ,						a_			)


							def    UpperCamelCase     (				self   :						int			)			->					Any:
										snake_case:    List[Any]      							=WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)
										snake_case:    Any      							=AutoProcessor.from_pretrained('hf-internal-testing/processor_with_lm'			)

										snake_case:    str      							=floats_list((3, 1_0_0_0)			)

										snake_case:    Optional[Any]      							=processor_wavaveca(a_      ,						return_tensors='np'			)
										snake_case:    Tuple      							=processor_auto(a_      ,						return_tensors='np'			)

										for key in input_wavaveca.keys():
													self.assertAlmostEqual(input_wavaveca[key].sum()      ,						input_auto[key].sum()      ,						delta=1E-2			)

										snake_case:    Union[str, Any]      							=self._get_dummy_logits()

										snake_case:    Tuple      							=processor_wavaveca.batch_decode(a_			)
										snake_case:    str      							=processor_auto.batch_decode(a_			)

										self.assertListEqual(decoded_wavaveca.text      ,						decoded_auto.text			)


							def    UpperCamelCase     (				self   :						Optional[Any]			)			->					Tuple:
										snake_case:    Optional[int]      							=self.get_feature_extractor()
										snake_case:    List[Any]      							=self.get_tokenizer()
										snake_case:    List[str]      							=self.get_decoder()

										snake_case:    Optional[int]      							=WavaVecaProcessorWithLM(tokenizer=a_      ,						feature_extractor=a_      ,						decoder=a_			)

										self.assertListEqual(
										    processor.model_input_names      ,						feature_extractor.model_input_names      ,						msg='`processor` and `feature_extractor` model input names do not match'      ,						)


							@staticmethod
							def    UpperCamelCase     (				a_   :						int      ,						a_   :						Optional[int]			)			->					Optional[int]:
										snake_case:    Optional[Any]      							=[d[key] for d in offsets]
										return retrieved_list


							def    UpperCamelCase     (				self   :						Dict			)			->					Any:
										snake_case:    Any      							=WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)
										snake_case:    Optional[int]      							=self._get_dummy_logits()[0]

										snake_case:    Tuple      							=processor.decode(a_      ,						output_word_offsets=a_			)
										# check Wav2Vec2CTCTokenizerOutput keys for word
										self.assertEqual(len(outputs.keys()			)      ,						4			)
										self.assertTrue('text' in outputs			)
										self.assertTrue('word_offsets' in outputs			)
										self.assertTrue(isinstance(a_      ,						a_			)			)

										self.assertEqual(' '.join(self.get_from_offsets(outputs['word_offsets']      ,						'word'			)			)      ,						outputs.text			)
										self.assertListEqual(self.get_from_offsets(outputs['word_offsets']      ,						'word'			)      ,						['<s>', '<s>', '</s>']			)
										self.assertListEqual(self.get_from_offsets(outputs['word_offsets']      ,						'start_offset'			)      ,						[0, 2, 4]			)
										self.assertListEqual(self.get_from_offsets(outputs['word_offsets']      ,						'end_offset'			)      ,						[1, 3, 5]			)


							def    UpperCamelCase     (				self   :						Dict			)			->					Optional[int]:
										snake_case:    Tuple      							=WavaVecaProcessorWithLM.from_pretrained('hf-internal-testing/processor_with_lm'			)
										snake_case:    int      							=self._get_dummy_logits()

										snake_case:    Optional[int]      							=processor.batch_decode(a_      ,						output_word_offsets=a_			)

										# check Wav2Vec2CTCTokenizerOutput keys for word
										self.assertEqual(len(outputs.keys()			)      ,						4			)
										self.assertTrue('text' in outputs			)
										self.assertTrue('word_offsets' in outputs			)
										self.assertTrue(isinstance(a_      ,						a_			)			)

										self.assertListEqual(
										    [' '.join(self.get_from_offsets(a_      ,						'word'			)			) for o in outputs['word_offsets']]      ,						outputs.text			)
										self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]      ,						'word'			)      ,						['<s>', '<s>', '</s>']			)
										self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]      ,						'start_offset'			)      ,						[0, 2, 4]			)
										self.assertListEqual(self.get_from_offsets(outputs['word_offsets'][0]      ,						'end_offset'			)      ,						[1, 3, 5]			)




							@slow
							@require_torch
							@require_torchaudio
							def    UpperCamelCase     (				self   :						List[str]			)			->					List[str]:
										import torch

										snake_case:    List[str]      							=load_dataset('common_voice'      ,						'en'      ,						split='train'      ,						streaming=a_			)
										snake_case:    Tuple      							=ds.cast_column('audio'      ,						datasets.Audio(sampling_rate=1_6_0_0_0			)			)
										snake_case:    List[Any]      							=iter(a_			)
										snake_case:    Optional[Any]      							=next(a_			)

										snake_case:    Tuple      							=AutoProcessor.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'			)
										snake_case:    Tuple      							=WavaVecaForCTC.from_pretrained('patrickvonplaten/wav2vec2-base-100h-with-lm'			)

										# compare to filename `common_voice_en_100038.mp3` of dataset viewer on https://huggingface.co/datasets/common_voice/viewer/en/train
										snake_case:    int      							=processor(sample['audio']['array']      ,						return_tensors='pt'			).input_values

										with torch.no_grad():
													snake_case:    Optional[int]      							=model(a_			).logits.cpu().numpy()

										snake_case:    int      							=processor.decode(logits[0]      ,						output_word_offsets=a_			)

										snake_case:    Tuple      							=model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
										snake_case:    Optional[int]      							=[
										    {
										        'start_time': d['start_offset'] * time_offset,
										        'end_time': d['end_offset'] * time_offset,
										        'word': d['word'],
										    }
										    for d in output['word_offsets']
										]

										snake_case:    Any      							='WHY DOES MILISANDRA LOOK LIKE SHE WANTS TO CONSUME JOHN SNOW ON THE RIVER AT THE WALL'

										# output words
										self.assertEqual(' '.join(self.get_from_offsets(a_      ,						'word'			)			)      ,						a_			)
										self.assertEqual(' '.join(self.get_from_offsets(a_      ,						'word'			)			)      ,						output.text			)

										# output times
										snake_case:    Tuple      							=torch.tensor(self.get_from_offsets(a_      ,						'start_time'			)			)
										snake_case:    Dict      							=torch.tensor(self.get_from_offsets(a_      ,						'end_time'			)			)

										# fmt: off
										snake_case:    Optional[Any]      							=torch.tensor([1.4_1_9_9, 1.6_5_9_9, 2.2_5_9_9, 3.0, 3.2_4, 3.5_9_9_9, 3.7_9_9_9, 4.0_9_9_9, 4.2_6, 4.9_4, 5.2_8, 5.6_5_9_9, 5.7_8, 5.9_4, 6.3_2, 6.5_3_9_9, 6.6_5_9_9]			)
										snake_case:    Optional[int]      							=torch.tensor([1.5_3_9_9, 1.8_9_9_9, 2.9, 3.1_6, 3.5_3_9_9, 3.7_2, 4.0_1_9_9, 4.1_7_9_9, 4.7_6, 5.1_5_9_9, 5.5_5_9_9, 5.6_9_9_9, 5.8_6, 6.1_9_9_9, 6.3_8, 6.6_1_9_9, 6.9_4]			)
										# fmt: on

										self.assertTrue(torch.allclose(a_      ,						a_      ,						atol=0.0_1			)			)
										self.assertTrue(torch.allclose(a_      ,						a_      ,						atol=0.0_1			)			)









'''simple docstring'''

import argparse
import json
import os
import re

import torch

from transformers import BloomConfig, BloomModel
from transformers.file_utils import CONFIG_NAME, WEIGHTS_NAME
from transformers.utils import logging


logging.set_verbosity_info()

a							   =				[
    'word_embeddings_layernorm.weight',
    'word_embeddings_layernorm.bias',
    'input_layernorm.weight',
    'input_layernorm.bias',
    'post_attention_layernorm.weight',
    'post_attention_layernorm.bias',
    'self_attention.dense.bias',
    'mlp.dense_4h_to_h.bias',
    'ln_f.weight',
    'ln_f.bias',
]

a							   =				[
    'mlp.dense_4h_to_h.weight',
    'self_attention.dense.weight',
]







def a_     (    __UpperCAmelCase				,							__UpperCAmelCase					)				->					Any:





			"""simple docstring"""







			snake_case:    Any      							={
			    'word_embeddings.weight': 'word_embeddings.weight',
			    'word_embeddings.norm.weight': 'word_embeddings_layernorm.weight',
			    'word_embeddings.norm.bias': 'word_embeddings_layernorm.bias',
			    'weight': 'ln_f.weight',
			    'bias': 'ln_f.bias',
			}

			if key in layer_rename_map:
						return layer_rename_map[key]

			# Handle transformer blocks
			snake_case:    Dict      							=int(re.match(R'.*layer_(\d*).*'				,							__UpperCAmelCase					)[1]					)
			layer_number -= 3
			return f'''h.{layer_number}.''' + key







def a_     (    __UpperCAmelCase					)				->					Union[str, Any]:





			"""simple docstring"""







			if dtype == torch.bool:
						return 1 / 8
			snake_case:    Optional[Any]      							=re.search(R'[^\d](\d+)$'				,							str(__UpperCAmelCase					)					)
			if bit_search is None:
						raise ValueError(f'''`dtype` is not a valid dtype: {dtype}.'''					)
			snake_case:    str      							=int(bit_search.groups()[0]					)
			return bit_size // 8







def a_     (    __UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase				,							__UpperCAmelCase					)				->					Tuple:





			"""simple docstring"""







			if bloom_config_file == "":
						snake_case:    Tuple      							=BloomConfig()
			else:
						snake_case:    List[Any]      							=BloomConfig.from_json_file(__UpperCAmelCase					)

			if shard_model:
						snake_case:    Any      							=os.listdir(__UpperCAmelCase					)
						snake_case:    Any      							=sorted(filter(lambda __UpperCAmelCase					: s.startswith('layer'					) and "model_00" in s				,							__UpperCAmelCase					)					)

						snake_case:    int      							={'weight_map': {}, 'metadata': {}}
						snake_case:    Dict      							=0

						snake_case:    List[str]      							=None

						snake_case:    Any      							=BloomConfig()

						for j, file in enumerate(__UpperCAmelCase					):
									print('Processing file: {}'.format(__UpperCAmelCase					)					)
									snake_case:    Any      							=None

									for i in range(__UpperCAmelCase					):
												# load all TP files
												snake_case:    Optional[int]      							=file.replace('model_00'				,							f'''model_0{i}'''					)
												snake_case:    List[str]      							=torch.load(os.path.join(__UpperCAmelCase				,							__UpperCAmelCase					)				,							map_location='cpu'					)

												# Rename keys in the transformers names
												snake_case:    int      							=list(temp.keys()					)
												for key in keys:
															snake_case:    List[Any]      							=temp.pop(__UpperCAmelCase					)

												if tensors is None:
															snake_case:    Tuple      							=temp
												else:
															for key in tensors.keys():
																		if any(key.endswith(__UpperCAmelCase					) for end in WEIGHTS_TO_AVERAGE_ENDSWITH					):
																					# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																					tensors[key] += temp[key]
																		else:
																					# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																					snake_case:    Optional[int]      							=1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN					) else 0
																					# We concatenate these weights accross TP ranks
																					snake_case:    str      							=torch.cat([tensors[key], temp[key]]				,							dim=__UpperCAmelCase					)

            # Divide by the number of TP the weights we want to average
									for key in tensors.keys():
												if any(key.endswith(__UpperCAmelCase					) for end in WEIGHTS_TO_AVERAGE_ENDSWITH					):
															snake_case:    Dict      							=tensors[key] / pretraining_tp
									torch.save(
									    __UpperCAmelCase				,							os.path.join(
									        __UpperCAmelCase				,							'pytorch_model_{}-of-{}.bin'.format(str(j + 1					).zfill(5					)				,							str(len(__UpperCAmelCase					)					).zfill(5					)					)				,							)				,							)

									for key in tensors.keys():
												snake_case:    List[str]      							=tensors[key]
												total_size += value.numel() * get_dtype_size(value.dtype					)
												if key not in index_dict["weight_map"]:
															snake_case:    List[Any]      							='pytorch_model_{}-of-{}.bin'.format(
															    str(j + 1					).zfill(5					)				,							str(len(__UpperCAmelCase					)					).zfill(5					)					)

						snake_case:    Tuple      							=BloomConfig()
						snake_case:    Union[str, Any]      							=pytorch_dump_folder_path + '/' + CONFIG_NAME
						snake_case:    Any      							=total_size
						with open(__UpperCAmelCase				,							'w'				,							encoding='utf-8'					) as f:
									f.write(config.to_json_string()					)
						with open(os.path.join(__UpperCAmelCase				,							WEIGHTS_NAME + '.index.json'					)				,							'w'				,							encoding='utf-8'					) as f:
									snake_case:    Tuple      							=json.dumps(__UpperCAmelCase				,							indent=2				,							sort_keys=__UpperCAmelCase					) + '\n'
									f.write(__UpperCAmelCase					)
			else:
						snake_case:    Optional[Any]      							=BloomModel(__UpperCAmelCase					)

						snake_case:    Any      							=os.listdir(__UpperCAmelCase					)
						snake_case:    Optional[int]      							=sorted(filter(lambda __UpperCAmelCase					: s.startswith('layer'					) and "model_00" in s				,							__UpperCAmelCase					)					)

						snake_case:    str      							=None
						for i, file in enumerate(__UpperCAmelCase					):
									snake_case:    List[str]      							=None
									for i in range(__UpperCAmelCase					):
												# load all TP files
												snake_case:    int      							=file.replace('model_00'				,							f'''model_0{i}'''					)
												snake_case:    Optional[Any]      							=torch.load(os.path.join(__UpperCAmelCase				,							__UpperCAmelCase					)				,							map_location='cpu'					)

												# Rename keys in the transformers names
												snake_case:    List[str]      							=list(temp.keys()					)
												for key in keys:
															snake_case:    int      							=temp.pop(__UpperCAmelCase					)

												if tensors is None:
															snake_case:    Any      							=temp
												else:
															for key in tensors.keys():
																		# We average (sum and then divide) some weights accross TP ranks (see https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/olruwase/sync_layer_norms/megatron/training.py#L425)
																		if any(key.endswith(__UpperCAmelCase					) for end in WEIGHTS_TO_AVERAGE_ENDSWITH					):
																					tensors[key] += temp[key]
																		else:
																					# Some weights are RowParallelLinear in Megatron-Deepspeed, others are ColumnParallel
																					snake_case:    List[Any]      							=1 if any(text in key for text in WEIGHTS_WITH_ROW_PARALLELISM_CONTAIN					) else 0
																					# We concatenate these weights accross TP ranks
																					snake_case:    Tuple      							=torch.cat([tensors[key], temp[key]]				,							dim=__UpperCAmelCase					)

            # Divide by the number of TP the weights we want to average
									for key in tensors.keys():
												if any(key.endswith(__UpperCAmelCase					) for end in WEIGHTS_TO_AVERAGE_ENDSWITH					):
															snake_case:    int      							=tensors[key] / pretraining_tp

									snake_case:    Optional[Any]      							=model.load_state_dict(__UpperCAmelCase				,							strict=__UpperCAmelCase					)
									assert not other_keys.unexpected_keys, f'''The keys {other_keys.unexpected_keys} are unexpected'''
									if missing_keys is None:
												snake_case:    Optional[Any]      							=set(other_keys.missing_keys					)
									else:
												snake_case:    Any      							=missing_keys.intersection(set(other_keys.missing_keys					)					)

						assert not missing_keys, f'''The keys {missing_keys} are missing'''

						# Save pytorch-model
						os.makedirs(__UpperCAmelCase				,							exist_ok=__UpperCAmelCase					)
						snake_case:    Optional[int]      							=pytorch_dump_folder_path + '/' + WEIGHTS_NAME
						snake_case:    Any      							=pytorch_dump_folder_path + '/' + CONFIG_NAME
						print(f'''Save PyTorch model to {pytorch_weights_dump_path} with dtype {config.torch_dtype}'''					)
						if config.torch_dtype is not None:
									snake_case:    Optional[Any]      							=model.to(config.torch_dtype					)
						torch.save(model.state_dict()				,							__UpperCAmelCase					)
						print(f'''Save configuration file to {pytorch_config_dump_path}'''					)
						with open(__UpperCAmelCase				,							'w'				,							encoding='utf-8'					) as f:
									f.write(config.to_json_string()					)


if __name__ == "__main__":
	a							   =				argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '--bloom_checkpoint_path',
	    default=None,
	    type=str,
	    required=True,
	    help='Path to the Megatron-LM checkpoint path.',
	)
	parser.add_argument(
	    '--pytorch_dump_folder_path', default=None, type=str, required=True, help='Path to the output PyTorch model.'
	)
	parser.add_argument(
	    '--bloom_config_file',
	    default='',
	    type=str,
	    help=(
	        'An optional config json file corresponding to the pre-trained model. \n'
	        'This specifies the model architecture.'
	    ),
	)
	parser.add_argument(
	    '--shard_model',
	    action='store_true',
	    help='An optional setting to shard the output model \nThis enables sharding the converted checkpoint',
	)
	parser.add_argument(
	    '--pretraining_tp',
	    default=4,
	    type=int,
	    help='Pretraining TP rank that has been used when training the model in Megatron-LM \n',
	)
	a							   =				parser.parse_args()
	convert_bloom_checkpoint_to_pytorch(
	    args.bloom_checkpoint_path,
	    args.bloom_config_file,
	    args.pytorch_dump_folder_path,
	    args.shard_model,
	    args.pretraining_tp,
	)




