



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        _snake_case    (self				):
						lowerCamelCase__ :      Union[str, Any]								=		"""| <pad> <unk> <s> </s> a b c d e f g h i j k""".split()
						lowerCamelCase__ :      int								=		dict(zip(__magic_name__							,    range(len(__magic_name__				)				)				)				)

						lowerCamelCase__ :      List[Any]								=		{
						    """unk_token""": """<unk>""",
						    """bos_token""": """<s>""",
						    """eos_token""": """</s>""",
						}
						lowerCamelCase__ :      Tuple								=		{
						    """feature_size""": 1,
						    """padding_value""": 0.0,
						    """sampling_rate""": 16000,
						    """return_attention_mask""": False,
						    """do_normalize""": True,
						}

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

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

						# load decoder from hub
						lowerCamelCase__ :      List[str]								=		"""hf-internal-testing/ngram-beam-search-decoder"""





			def        _snake_case    (self							,    **__magic_name__				):
						lowerCamelCase__ :      Union[str, Any]								=		self.add_kwargs_tokens_map.copy()
						kwargs.update(__magic_name__				)
						return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname							,    **__magic_name__				)





			def        _snake_case    (self							,    **__magic_name__				):
						return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname							,    **__magic_name__				)





			def        _snake_case    (self							,    **__magic_name__				):
						return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name							,    **__magic_name__				)





			def        _snake_case    (self				):
						shutil.rmtree(self.tmpdirname				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		self.get_tokenizer()
						lowerCamelCase__ :      Optional[int]								=		self.get_feature_extractor()
						lowerCamelCase__ :      int								=		self.get_decoder()

						lowerCamelCase__ :      List[str]								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

						processor.save_pretrained(self.tmpdirname				)
						lowerCamelCase__ :      Dict								=		WavaVecaProcessorWithLM.from_pretrained(self.tmpdirname				)

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

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

						# 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							,    __magic_name__				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		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
						lowerCamelCase__ :      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        _snake_case    (self				):
						lowerCamelCase__ :      int								=		self.get_tokenizer()
						# add token to trigger raise
						tokenizer.add_tokens(["""xx"""]				)
						with self.assertRaisesRegex(__magic_name__							,    """include"""				):
									WavaVecaProcessorWithLM(
									    tokenizer=__magic_name__							,    feature_extractor=self.get_feature_extractor()							,    decoder=self.get_decoder()				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Optional[Any]								=		self.get_feature_extractor()
						lowerCamelCase__ :      Union[str, Any]								=		self.get_tokenizer()
						lowerCamelCase__ :      Any								=		self.get_decoder()

						lowerCamelCase__ :      Any								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

						lowerCamelCase__ :      Tuple								=		floats_list((3, 1000)				)

						lowerCamelCase__ :      Any								=		feature_extractor(__magic_name__							,    return_tensors="""np"""				)
						lowerCamelCase__ :      int								=		processor(__magic_name__							,    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        _snake_case    (self				):
						lowerCamelCase__ :      Optional[int]								=		self.get_feature_extractor()
						lowerCamelCase__ :      Union[str, Any]								=		self.get_tokenizer()
						lowerCamelCase__ :      int								=		self.get_decoder()

						lowerCamelCase__ :      List[str]								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

						lowerCamelCase__ :      Optional[int]								=		"""This is a test string"""

						lowerCamelCase__ :      List[str]								=		processor(text=__magic_name__				)

						lowerCamelCase__ :      Any								=		tokenizer(__magic_name__				)

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





			def        _snake_case    (self							,    __magic_name__=(2, 10, 16)							,    __magic_name__=77				):
						np.random.seed(__magic_name__				)
						return np.random.rand(*__magic_name__				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      str								=		self.get_feature_extractor()
						lowerCamelCase__ :      Any								=		self.get_tokenizer()
						lowerCamelCase__ :      Optional[Any]								=		self.get_decoder()

						lowerCamelCase__ :      Optional[int]								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

						lowerCamelCase__ :      Tuple								=		self._get_dummy_logits(shape=(10, 16)							,    seed=13				)

						lowerCamelCase__ :      Optional[Any]								=		processor.decode(__magic_name__				)

						lowerCamelCase__ :      List[Any]								=		decoder.decode_beams(__magic_name__				)[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        _snake_case    (self							,    __magic_name__				):
						lowerCamelCase__ :      List[str]								=		self.get_feature_extractor()
						lowerCamelCase__ :      List[str]								=		self.get_tokenizer()
						lowerCamelCase__ :      str								=		self.get_decoder()

						lowerCamelCase__ :      str								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

						lowerCamelCase__ :      str								=		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:
									lowerCamelCase__ :      List[str]								=		processor.batch_decode(__magic_name__				)
						else:
									with get_context(__magic_name__				).Pool() as pool:
												lowerCamelCase__ :      Tuple								=		processor.batch_decode(__magic_name__							,    __magic_name__				)

						lowerCamelCase__ :      str								=		list(__magic_name__				)

						with get_context("""fork"""				).Pool() as p:
									lowerCamelCase__ :      List[str]								=		decoder.decode_beams_batch(__magic_name__							,    __magic_name__				)

						lowerCamelCase__							,lowerCamelCase__							,lowerCamelCase__ :      Dict								=		[], [], []
						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(__magic_name__							,    decoded_processor.text				)
						self.assertListEqual(["""<s> <s> </s>""", """<s> <s> <s>"""]							,    decoded_processor.text				)
						self.assertListEqual(__magic_name__							,    decoded_processor.logit_score				)
						self.assertListEqual(__magic_name__							,    decoded_processor.lm_score				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Any								=		self.get_feature_extractor()
						lowerCamelCase__ :      Optional[int]								=		self.get_tokenizer()
						lowerCamelCase__ :      Any								=		self.get_decoder()

						lowerCamelCase__ :      List[Any]								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

						lowerCamelCase__ :      int								=		self._get_dummy_logits()

						lowerCamelCase__ :      Optional[Any]								=		15
						lowerCamelCase__ :      Union[str, Any]								=		-20.0
						lowerCamelCase__ :      List[Any]								=		-4.0

						lowerCamelCase__ :      List[str]								=		processor.batch_decode(
						    __magic_name__							,    beam_width=__magic_name__							,    beam_prune_logp=__magic_name__							,    token_min_logp=__magic_name__							,    )
						lowerCamelCase__ :      str								=		decoded_processor_out.text

						lowerCamelCase__ :      int								=		list(__magic_name__				)

						with get_context("""fork"""				).Pool() as pool:
									lowerCamelCase__ :      List[Any]								=		decoder.decode_beams_batch(
									    __magic_name__							,    __magic_name__							,    beam_width=__magic_name__							,    beam_prune_logp=__magic_name__							,    token_min_logp=__magic_name__							,    )

						lowerCamelCase__ :      str								=		[d[0][0] for d in decoded_decoder_out]
						lowerCamelCase__ :      List[Any]								=		[d[0][2] for d in decoded_decoder_out]
						lowerCamelCase__ :      Optional[Any]								=		[d[0][3] for d in decoded_decoder_out]

						self.assertListEqual(__magic_name__							,    __magic_name__				)
						self.assertListEqual(["""</s> <s> <s>""", """<s> <s> <s>"""]							,    __magic_name__				)

						self.assertTrue(np.array_equal(__magic_name__							,    decoded_processor_out.logit_score				)				)
						self.assertTrue(np.allclose([-20.0_54, -18.4_47]							,    __magic_name__							,    atol=1E-3				)				)

						self.assertTrue(np.array_equal(__magic_name__							,    decoded_processor_out.lm_score				)				)
						self.assertTrue(np.allclose([-15.5_54, -13.94_74]							,    __magic_name__							,    atol=1E-3				)				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Union[str, Any]								=		self.get_feature_extractor()
						lowerCamelCase__ :      Dict								=		self.get_tokenizer()
						lowerCamelCase__ :      List[str]								=		self.get_decoder()

						lowerCamelCase__ :      Union[str, Any]								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

						lowerCamelCase__ :      Any								=		self._get_dummy_logits()

						lowerCamelCase__ :      str								=		2.0
						lowerCamelCase__ :      Tuple								=		5.0
						lowerCamelCase__ :      Dict								=		-20.0
						lowerCamelCase__ :      List[Any]								=		True

						lowerCamelCase__ :      Union[str, Any]								=		processor.batch_decode(
						    __magic_name__							,    alpha=__magic_name__							,    beta=__magic_name__							,    unk_score_offset=__magic_name__							,    lm_score_boundary=__magic_name__							,    )
						lowerCamelCase__ :      List[str]								=		decoded_processor_out.text

						lowerCamelCase__ :      Optional[Any]								=		list(__magic_name__				)
						decoder.reset_params(
						    alpha=__magic_name__							,    beta=__magic_name__							,    unk_score_offset=__magic_name__							,    lm_score_boundary=__magic_name__							,    )

						with get_context("""fork"""				).Pool() as pool:
									lowerCamelCase__ :      int								=		decoder.decode_beams_batch(
									    __magic_name__							,    __magic_name__							,    )

						lowerCamelCase__ :      Optional[Any]								=		[d[0][0] for d in decoded_decoder_out]

						self.assertListEqual(__magic_name__							,    __magic_name__				)
						self.assertListEqual(["""<s> </s> <s> </s> </s>""", """</s> </s> <s> </s> </s>"""]							,    __magic_name__				)
						lowerCamelCase__ :      Optional[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							,    -20.0				)
						self.assertEqual(lm_model.score_boundary							,    __magic_name__				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Union[str, Any]								=		WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)

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

						lowerCamelCase__ :      Any								=		os.listdir(__magic_name__				)
						lowerCamelCase__ :      List[Any]								=		["""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(__magic_name__							,    __magic_name__				)





			def        _snake_case    (self				):
						lowerCamelCase__ :      Tuple								=		snapshot_download("""hf-internal-testing/processor_with_lm"""				)

						lowerCamelCase__ :      Tuple								=		WavaVecaProcessorWithLM.from_pretrained(__magic_name__				)

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

						lowerCamelCase__ :      Any								=		os.listdir(__magic_name__				)
						lowerCamelCase__ :      List[str]								=		os.listdir(__magic_name__				)

						local_decoder_files.sort()
						expected_decoder_files.sort()

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





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[str]								=		WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)
						lowerCamelCase__ :      Any								=		AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""				)

						lowerCamelCase__ :      Union[str, Any]								=		floats_list((3, 1000)				)

						lowerCamelCase__ :      Tuple								=		processor_wavaveca(__magic_name__							,    return_tensors="""np"""				)
						lowerCamelCase__ :      Tuple								=		processor_auto(__magic_name__							,    return_tensors="""np"""				)

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

						lowerCamelCase__ :      str								=		self._get_dummy_logits()

						lowerCamelCase__ :      int								=		processor_wavaveca.batch_decode(__magic_name__				)
						lowerCamelCase__ :      Tuple								=		processor_auto.batch_decode(__magic_name__				)

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





			def        _snake_case    (self				):
						lowerCamelCase__ :      List[Any]								=		self.get_feature_extractor()
						lowerCamelCase__ :      Dict								=		self.get_tokenizer()
						lowerCamelCase__ :      Optional[Any]								=		self.get_decoder()

						lowerCamelCase__ :      Dict								=		WavaVecaProcessorWithLM(tokenizer=__magic_name__							,    feature_extractor=__magic_name__							,    decoder=__magic_name__				)

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





			@staticmethod
			def        _snake_case    (__magic_name__							,    __magic_name__				):
						lowerCamelCase__ :      List[Any]								=		[d[key] for d in offsets]
						return retrieved_list





			def        _snake_case    (self				):
						lowerCamelCase__ :      int								=		WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)
						lowerCamelCase__ :      int								=		self._get_dummy_logits()[0]

						lowerCamelCase__ :      List[Any]								=		processor.decode(__magic_name__							,    output_word_offsets=__magic_name__				)
						# 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(__magic_name__							,    __magic_name__				)				)

						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        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""				)
						lowerCamelCase__ :      Tuple								=		self._get_dummy_logits()

						lowerCamelCase__ :      Optional[int]								=		processor.batch_decode(__magic_name__							,    output_word_offsets=__magic_name__				)

						# 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(__magic_name__							,    __magic_name__				)				)

						self.assertListEqual(
						    [""" """.join(self.get_from_offsets(__magic_name__							,    """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        _snake_case    (self				):
						import torch

						lowerCamelCase__ :      Tuple								=		load_dataset("""common_voice"""							,    """en"""							,    split="""train"""							,    streaming=__magic_name__				)
						lowerCamelCase__ :      str								=		ds.cast_column("""audio"""							,    datasets.Audio(sampling_rate=16000				)				)
						lowerCamelCase__ :      Optional[int]								=		iter(__magic_name__				)
						lowerCamelCase__ :      Tuple								=		next(__magic_name__				)

						lowerCamelCase__ :      Any								=		AutoProcessor.from_pretrained("""patrickvonplaten/wav2vec2-base-100h-with-lm"""				)
						lowerCamelCase__ :      str								=		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
						lowerCamelCase__ :      Tuple								=		processor(sample["""audio"""]["""array"""]							,    return_tensors="""pt"""				).input_values

						with torch.no_grad():
									lowerCamelCase__ :      str								=		model(__magic_name__				).logits.cpu().numpy()

						lowerCamelCase__ :      List[str]								=		processor.decode(logits[0]							,    output_word_offsets=__magic_name__				)

						lowerCamelCase__ :      str								=		model.config.inputs_to_logits_ratio / processor.feature_extractor.sampling_rate
						lowerCamelCase__ :      Union[str, Any]								=		[
						    {
						        """start_time""": d["""start_offset"""] * time_offset,
						        """end_time""": d["""end_offset"""] * time_offset,
						        """word""": d["""word"""],
						    }
						    for d in output["""word_offsets"""]
						]

						lowerCamelCase__ :      str								=		"""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(__magic_name__							,    """word"""				)				)							,    __magic_name__				)
						self.assertEqual(""" """.join(self.get_from_offsets(__magic_name__							,    """word"""				)				)							,    output.text				)

						# output times
						lowerCamelCase__ :      Any								=		torch.tensor(self.get_from_offsets(__magic_name__							,    """start_time"""				)				)
						lowerCamelCase__ :      Dict								=		torch.tensor(self.get_from_offsets(__magic_name__							,    """end_time"""				)				)

						# fmt: off
						lowerCamelCase__ :      Any								=		torch.tensor([1.41_99, 1.65_99, 2.25_99, 3.0, 3.24, 3.59_99, 3.79_99, 4.09_99, 4.26, 4.94, 5.28, 5.65_99, 5.78, 5.94, 6.32, 6.53_99, 6.65_99]				)
						lowerCamelCase__ :      str								=		torch.tensor([1.53_99, 1.89_99, 2.9, 3.16, 3.53_99, 3.72, 4.01_99, 4.17_99, 4.76, 5.15_99, 5.55_99, 5.69_99, 5.86, 6.19_99, 6.38, 6.61_99, 6.94]				)
						# fmt: on

						self.assertTrue(torch.allclose(__magic_name__							,    __magic_name__							,    atol=0.01				)				)
						self.assertTrue(torch.allclose(__magic_name__							,    __magic_name__							,    atol=0.01				)				)








import argparse
import glob
import logging
import os
from argparse import Namespace
from importlib import import_module

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from seqeval.metrics import accuracy_score, fa_score, precision_score, recall_score
from torch.nn import CrossEntropyLoss
from torch.utils.data import DataLoader, TensorDataset
from utils_ner import TokenClassificationTask


_lowercase     							=						logging.getLogger(__name__)
class 				__A					( A_      ):
			UpperCamelCase					:Optional[int]     =			'''token-classification'''





			def __init__(self							,    __magic_name__				):
						if type(__magic_name__				) == dict:
									lowerCamelCase__ :      Any								=		Namespace(**__magic_name__				)
						lowerCamelCase__ :      str								=		import_module("""tasks"""				)
						try:
									lowerCamelCase__ :      Optional[Any]								=		getattr(__magic_name__							,    hparams.task_type				)
									lowerCamelCase__ :      TokenClassificationTask								=		token_classification_task_clazz()
						except AttributeError:
									raise ValueError(
									    f"Task {hparams.task_type} needs to be defined as a TokenClassificationTask subclass in {module}. "
									    f"Available tasks classes are: {TokenClassificationTask.__subclasses__()}"				)
						lowerCamelCase__ :      Any								=		self.token_classification_task.get_labels(hparams.labels				)
						lowerCamelCase__ :      Tuple								=		CrossEntropyLoss().ignore_index
						super().__init__(__magic_name__							,    len(self.labels				)							,    self.mode				)





			def        _snake_case    (self							,    **__magic_name__				):
						return self.model(**__magic_name__				)





			def        _snake_case    (self							,    __magic_name__							,    __magic_name__				):
						lowerCamelCase__ :      Tuple								=		{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
						if self.config.model_type != "distilbert":
									lowerCamelCase__ :      Union[str, Any]								=		(
									    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
									)  # XLM and RoBERTa don"t use token_type_ids

						lowerCamelCase__ :      List[str]								=		self(**__magic_name__				)
						lowerCamelCase__ :      Dict								=		outputs[0]
						# tensorboard_logs = {"loss": loss, "rate": self.lr_scheduler.get_last_lr()[-1]}
						return {"loss": loss}





			def        _snake_case    (self				):
						lowerCamelCase__ :      Dict								=		self.hparams
						for mode in ["train", "dev", "test"]:
									lowerCamelCase__ :      List[str]								=		self._feature_file(__magic_name__				)
									if os.path.exists(__magic_name__				) and not args.overwrite_cache:
												logger.info("""Loading features from cached file %s"""							,    __magic_name__				)
												lowerCamelCase__ :      Union[str, Any]								=		torch.load(__magic_name__				)
									else:
												logger.info("""Creating features from dataset file at %s"""							,    args.data_dir				)
												lowerCamelCase__ :      int								=		self.token_classification_task.read_examples_from_file(args.data_dir							,    __magic_name__				)
												lowerCamelCase__ :      Tuple								=		self.token_classification_task.convert_examples_to_features(
												    __magic_name__							,    self.labels							,    args.max_seq_length							,    self.tokenizer							,    cls_token_at_end=bool(self.config.model_type in ["""xlnet"""]				)							,    cls_token=self.tokenizer.cls_token							,    cls_token_segment_id=2 if self.config.model_type in ["""xlnet"""] else 0							,    sep_token=self.tokenizer.sep_token							,    sep_token_extra=__magic_name__							,    pad_on_left=bool(self.config.model_type in ["""xlnet"""]				)							,    pad_token=self.tokenizer.pad_token_id							,    pad_token_segment_id=self.tokenizer.pad_token_type_id							,    pad_token_label_id=self.pad_token_label_id							,    )
												logger.info("""Saving features into cached file %s"""							,    __magic_name__				)
												torch.save(__magic_name__							,    __magic_name__				)





			def        _snake_case    (self							,    __magic_name__							,    __magic_name__							,    __magic_name__ = False				):
						lowerCamelCase__ :      Any								=		self._feature_file(__magic_name__				)
						logger.info("""Loading features from cached file %s"""							,    __magic_name__				)
						lowerCamelCase__ :      Optional[Any]								=		torch.load(__magic_name__				)
						lowerCamelCase__ :      Tuple								=		torch.tensor([f.input_ids for f in features]							,    dtype=torch.long				)
						lowerCamelCase__ :      Optional[Any]								=		torch.tensor([f.attention_mask for f in features]							,    dtype=torch.long				)
						if features[0].token_type_ids is not None:
									lowerCamelCase__ :      str								=		torch.tensor([f.token_type_ids for f in features]							,    dtype=torch.long				)
						else:
									lowerCamelCase__ :      int								=		torch.tensor([0 for f in features]							,    dtype=torch.long				)
									# HACK(we will not use this anymore soon)
						lowerCamelCase__ :      Any								=		torch.tensor([f.label_ids for f in features]							,    dtype=torch.long				)
						return DataLoader(
						    TensorDataset(__magic_name__							,    __magic_name__							,    __magic_name__							,    __magic_name__				)							,    batch_size=__magic_name__				)





			def        _snake_case    (self							,    __magic_name__							,    __magic_name__				):
						"""Compute validation""" ""
						lowerCamelCase__ :      Optional[int]								=		{"""input_ids""": batch[0], """attention_mask""": batch[1], """labels""": batch[3]}
						if self.config.model_type != "distilbert":
									lowerCamelCase__ :      Tuple								=		(
									    batch[2] if self.config.model_type in ["""bert""", """xlnet"""] else None
									)  # XLM and RoBERTa don"t use token_type_ids
						lowerCamelCase__ :      str								=		self(**__magic_name__				)
						lowerCamelCase__							,lowerCamelCase__ :      List[Any]								=		outputs[:2]
						lowerCamelCase__ :      List[Any]								=		logits.detach().cpu().numpy()
						lowerCamelCase__ :      Dict								=		inputs["""labels"""].detach().cpu().numpy()
						return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}





			def        _snake_case    (self							,    __magic_name__				):
						lowerCamelCase__ :      List[str]								=		torch.stack([x["""val_loss"""] for x in outputs]				).mean()
						lowerCamelCase__ :      Optional[int]								=		np.concatenate([x["""pred"""] for x in outputs]							,    axis=0				)
						lowerCamelCase__ :      List[str]								=		np.argmax(__magic_name__							,    axis=2				)
						lowerCamelCase__ :      Optional[Any]								=		np.concatenate([x["""target"""] for x in outputs]							,    axis=0				)

						lowerCamelCase__ :      Optional[int]								=		dict(enumerate(self.labels				)				)
						lowerCamelCase__ :      List[str]								=		[[] for _ in range(out_label_ids.shape[0]				)]
						lowerCamelCase__ :      Any								=		[[] for _ in range(out_label_ids.shape[0]				)]

						for i in range(out_label_ids.shape[0]				):
									for j in range(out_label_ids.shape[1]				):
												if out_label_ids[i, j] != self.pad_token_label_id:
															out_label_list[i].append(label_map[out_label_ids[i][j]]				)
															preds_list[i].append(label_map[preds[i][j]]				)

						lowerCamelCase__ :      Tuple								=		{
						    """val_loss""": val_loss_mean,
						    """accuracy_score""": accuracy_score(__magic_name__							,    __magic_name__				),
						    """precision""": precision_score(__magic_name__							,    __magic_name__				),
						    """recall""": recall_score(__magic_name__							,    __magic_name__				),
						    """f1""": fa_score(__magic_name__							,    __magic_name__				),
						}

						lowerCamelCase__ :      Dict								=		dict(results.items()				)
						lowerCamelCase__ :      str								=		results
						return ret, preds_list, out_label_list





			def        _snake_case    (self							,    __magic_name__				):
						# when stable
						lowerCamelCase__							,lowerCamelCase__							,lowerCamelCase__ :      str								=		self._eval_end(__magic_name__				)
						lowerCamelCase__ :      Union[str, Any]								=		ret["""log"""]
						return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}





			def        _snake_case    (self							,    __magic_name__				):
						# updating to test_epoch_end instead of deprecated test_end
						lowerCamelCase__							,lowerCamelCase__							,lowerCamelCase__ :      Union[str, Any]								=		self._eval_end(__magic_name__				)

						# Converting to the dict required by pl
						# https://github.com/PyTorchLightning/pytorch-lightning/blob/master/\
						# pytorch_lightning/trainer/logging.py#L139
						lowerCamelCase__ :      List[Any]								=		ret["""log"""]
						# `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
						return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}





			@staticmethod
			def        _snake_case    (__magic_name__							,    __magic_name__				):
						# Add NER specific options
						BaseTransformer.add_model_specific_args(__magic_name__							,    __magic_name__				)
						parser.add_argument(
						    """--task_type"""							,    default="""NER"""							,    type=__magic_name__							,    help="""Task type to fine tune in training (e.g. NER, POS, etc)"""				)
						parser.add_argument(
						    """--max_seq_length"""							,    default=128							,    type=__magic_name__							,    help=(
						        """The maximum total input sequence length after tokenization. Sequences longer """
						        """than this will be truncated, sequences shorter will be padded."""
						    )							,    )

						parser.add_argument(
						    """--labels"""							,    default=""""""							,    type=__magic_name__							,    help="""Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."""							,    )
						parser.add_argument(
						    """--gpus"""							,    default=0							,    type=__magic_name__							,    help="""The number of GPUs allocated for this, it is by default 0 meaning none"""							,    )

						parser.add_argument(
						    """--overwrite_cache"""							,    action="""store_true"""							,    help="""Overwrite the cached training and evaluation sets"""				)

						return parser


if __name__ == "__main__":
				_lowercase     							=						argparse.ArgumentParser()
				add_generic_args(parser, os.getcwd())
				_lowercase     							=						NERTransformer.add_model_specific_args(parser, os.getcwd())
				_lowercase     							=						parser.parse_args()
				_lowercase     							=						NERTransformer(args)
				_lowercase     							=						generic_train(model, args)

				if args.do_predict:
								# See https://github.com/huggingface/transformers/issues/3159
								# pl use this default format to create a checkpoint:
								# https://github.com/PyTorchLightning/pytorch-lightning/blob/master\
								# /pytorch_lightning/callbacks/model_checkpoint.py#L322
								_lowercase     							=						sorted(glob.glob(os.path.join(args.output_dir, '''checkpoint-epoch=*.ckpt'''), recursive=True))
								_lowercase     							=						model.load_from_checkpoint(checkpoints[-1])
								trainer.test(model)



