




"""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        __magic_name__				(			unittest.TestCase							):






				'''simple docstring'''





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   """| <pad> <unk> <s> </s> a b c d e f g h i j k""".split()
							lowerCamelCase      =   dict(zip(_a						,					range(len(_a					)					)					)					)

							lowerCamelCase      =   {
							    """unk_token""": """<unk>""",
							    """bos_token""": """<s>""",
							    """eos_token""": """</s>""",
							}
							lowerCamelCase      =   {
							    """feature_size""": 1,
							    """padding_value""": 0.0,
							    """sampling_rate""": 16_000,
							    """return_attention_mask""": False,
							    """do_normalize""": True,
							}

							lowerCamelCase      =   tempfile.mkdtemp()
							lowerCamelCase      =   os.path.join(self.tmpdirname						,					VOCAB_FILES_NAMES["""vocab_file"""]					)
							lowerCamelCase      =   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
							lowerCamelCase      =   """hf-internal-testing/ngram-beam-search-decoder"""





				def   _lowerCAmelCase							(			self						,					**_a					):





							"""simple docstring"""
							lowerCamelCase      =   self.add_kwargs_tokens_map.copy()
							kwargs.update(_a					)
							return WavaVecaCTCTokenizer.from_pretrained(self.tmpdirname						,					**_a					)





				def   _lowerCAmelCase							(			self						,					**_a					):





							"""simple docstring"""
							return WavaVecaFeatureExtractor.from_pretrained(self.tmpdirname						,					**_a					)





				def   _lowerCAmelCase							(			self						,					**_a					):





							"""simple docstring"""
							return BeamSearchDecoderCTC.load_from_hf_hub(self.decoder_name						,					**_a					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							shutil.rmtree(self.tmpdirname					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   WavaVecaProcessorWithLM(tokenizer=_a						,					feature_extractor=_a						,					decoder=_a					)

							processor.save_pretrained(self.tmpdirname					)
							lowerCamelCase      =   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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   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      =   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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   WavaVecaProcessorWithLM(tokenizer=_a						,					feature_extractor=_a						,					decoder=_a					)

							lowerCamelCase      =   floats_list((3, 1_000)					)

							lowerCamelCase      =   feature_extractor(_a						,					return_tensors="""np"""					)
							lowerCamelCase      =   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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   WavaVecaProcessorWithLM(tokenizer=_a						,					feature_extractor=_a						,					decoder=_a					)

							lowerCamelCase      =   """This is a test string"""

							lowerCamelCase      =   processor(text=_a					)

							lowerCamelCase      =   tokenizer(_a					)

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





				def   _lowerCAmelCase							(			self						,					_a=(2, 10, 16)						,					_a=77					):





							"""simple docstring"""
							np.random.seed(_a					)
							return np.random.rand(*_a					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   WavaVecaProcessorWithLM(tokenizer=_a						,					feature_extractor=_a						,					decoder=_a					)

							lowerCamelCase      =   self._get_dummy_logits(shape=(10, 16)						,					seed=13					)

							lowerCamelCase      =   processor.decode(_a					)

							lowerCamelCase      =   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   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   WavaVecaProcessorWithLM(tokenizer=_a						,					feature_extractor=_a						,					decoder=_a					)

							lowerCamelCase      =   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      =   processor.batch_decode(_a					)
							else:
										with get_context(_a					).Pool() as pool:
													lowerCamelCase      =   processor.batch_decode(_a						,					_a					)

							lowerCamelCase      =   list(_a					)

							with get_context("""fork"""					).Pool() as p:
										lowerCamelCase      =   decoder.decode_beams_batch(_a						,					_a					)

							lowerCamelCase ,  lowerCamelCase ,  lowerCamelCase      =   [], [], []
							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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   WavaVecaProcessorWithLM(tokenizer=_a						,					feature_extractor=_a						,					decoder=_a					)

							lowerCamelCase      =   self._get_dummy_logits()

							lowerCamelCase      =   15
							lowerCamelCase      =   -20.0
							lowerCamelCase      =   -4.0

							lowerCamelCase      =   processor.batch_decode(
							    _a						,					beam_width=_a						,					beam_prune_logp=_a						,					token_min_logp=_a						,					)
							lowerCamelCase      =   decoded_processor_out.text

							lowerCamelCase      =   list(_a					)

							with get_context("""fork"""					).Pool() as pool:
										lowerCamelCase      =   decoder.decode_beams_batch(
										    _a						,					_a						,					beam_width=_a						,					beam_prune_logp=_a						,					token_min_logp=_a						,					)

							lowerCamelCase      =   [d[0][0] for d in decoded_decoder_out]
							lowerCamelCase      =   [d[0][2] for d in decoded_decoder_out]
							lowerCamelCase      =   [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([-20.054, -18.447]						,					_a						,					atol=1e-3					)					)

							self.assertTrue(np.array_equal(_a						,					decoded_processor_out.lm_score					)					)
							self.assertTrue(np.allclose([-15.554, -13.9_474]						,					_a						,					atol=1e-3					)					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   WavaVecaProcessorWithLM(tokenizer=_a						,					feature_extractor=_a						,					decoder=_a					)

							lowerCamelCase      =   self._get_dummy_logits()

							lowerCamelCase      =   2.0
							lowerCamelCase      =   5.0
							lowerCamelCase      =   -20.0
							lowerCamelCase      =   True

							lowerCamelCase      =   processor.batch_decode(
							    _a						,					alpha=_a						,					beta=_a						,					unk_score_offset=_a						,					lm_score_boundary=_a						,					)
							lowerCamelCase      =   decoded_processor_out.text

							lowerCamelCase      =   list(_a					)
							decoder.reset_params(
							    alpha=_a						,					beta=_a						,					unk_score_offset=_a						,					lm_score_boundary=_a						,					)

							with get_context("""fork"""					).Pool() as pool:
										lowerCamelCase      =   decoder.decode_beams_batch(
										    _a						,					_a						,					)

							lowerCamelCase      =   [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					)
							lowerCamelCase      =   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						,					_a					)





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""					)

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

							lowerCamelCase      =   os.listdir(_a					)
							lowerCamelCase      =   ["""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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   snapshot_download("""hf-internal-testing/processor_with_lm"""					)

							lowerCamelCase      =   WavaVecaProcessorWithLM.from_pretrained(_a					)

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

							lowerCamelCase      =   os.listdir(_a					)
							lowerCamelCase      =   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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""					)
							lowerCamelCase      =   AutoProcessor.from_pretrained("""hf-internal-testing/processor_with_lm"""					)

							lowerCamelCase      =   floats_list((3, 1_000)					)

							lowerCamelCase      =   processor_wavaveca(_a						,					return_tensors="""np"""					)
							lowerCamelCase      =   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					)

							lowerCamelCase      =   self._get_dummy_logits()

							lowerCamelCase      =   processor_wavaveca.batch_decode(_a					)
							lowerCamelCase      =   processor_auto.batch_decode(_a					)

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





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   self.get_feature_extractor()
							lowerCamelCase      =   self.get_tokenizer()
							lowerCamelCase      =   self.get_decoder()

							lowerCamelCase      =   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   _lowerCAmelCase							(			_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   [d[key] for d in offsets]
							return retrieved_list





				def   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""					)
							lowerCamelCase      =   self._get_dummy_logits()[0]

							lowerCamelCase      =   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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							lowerCamelCase      =   WavaVecaProcessorWithLM.from_pretrained("""hf-internal-testing/processor_with_lm"""					)
							lowerCamelCase      =   self._get_dummy_logits()

							lowerCamelCase      =   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   _lowerCAmelCase							(			self					):





							"""simple docstring"""
							import torch

							lowerCamelCase      =   load_dataset("""common_voice"""						,					"""en"""						,					split="""train"""						,					streaming=_a					)
							lowerCamelCase      =   ds.cast_column("""audio"""						,					datasets.Audio(sampling_rate=16_000					)					)
							lowerCamelCase      =   iter(_a					)
							lowerCamelCase      =   next(_a					)

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

							with torch.no_grad():
										lowerCamelCase      =   model(_a					).logits.cpu().numpy()

							lowerCamelCase      =   processor.decode(logits[0]						,					output_word_offsets=_a					)

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

							lowerCamelCase      =   """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
							lowerCamelCase      =   torch.tensor(self.get_from_offsets(_a						,					"""start_time"""					)					)
							lowerCamelCase      =   torch.tensor(self.get_from_offsets(_a						,					"""end_time"""					)					)

							# fmt: off
							lowerCamelCase      =   torch.tensor([1.4_199, 1.6_599, 2.2_599, 3.0, 3.24, 3.5_999, 3.7_999, 4.0_999, 4.26, 4.94, 5.28, 5.6_599, 5.78, 5.94, 6.32, 6.5_399, 6.6_599]					)
							lowerCamelCase      =   torch.tensor([1.5_399, 1.8_999, 2.9, 3.16, 3.5_399, 3.72, 4.0_199, 4.1_799, 4.76, 5.1_599, 5.5_599, 5.6_999, 5.86, 6.1_999, 6.38, 6.6_199, 6.94]					)
							# fmt: on

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









"""simple docstring"""




from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin






class        __magic_name__				(			UpperCAmelCase__		,   UpperCAmelCase__		,   UpperCAmelCase__							):






				'''simple docstring'''
				__UpperCamelCase         =		[R"h\.\d+\.attn\.bias", R"h\.\d+\.attn\.masked_bias"]





				@register_to_config
				def __init__(			self						,					_a						,					_a						,					_a = None						,					_a = 50_257						,					_a = 1_024						,					_a = 768						,					_a = 12						,					_a = 12						,					_a = None						,					_a = "gelu_new"						,					_a = 0.1						,					_a = 0.1						,					_a = 0.1						,					_a = 1e-5						,					_a = 0.02						,					_a = True						,					_a = True						,					_a = False						,					_a = False						,					):





							"""simple docstring"""
							super().__init__()

							lowerCamelCase      =   prefix_length

							if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
										raise ValueError(
										    f'`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'
										    f' `n_embd`: {n_embd} are not equal.'					)

							lowerCamelCase      =   prefix_inner_dim
							lowerCamelCase      =   prefix_hidden_dim

							lowerCamelCase      =   (
							    nn.Linear(self.prefix_inner_dim						,					self.prefix_hidden_dim					)
							    if self.prefix_hidden_dim is not None
							    else nn.Identity()
							)
							lowerCamelCase      =   (
							    nn.Linear(self.prefix_hidden_dim						,					_a					) if self.prefix_hidden_dim is not None else nn.Identity()
							)

							lowerCamelCase      =   GPTaConfig(
							    vocab_size=_a						,					n_positions=_a						,					n_embd=_a						,					n_layer=_a						,					n_head=_a						,					n_inner=_a						,					activation_function=_a						,					resid_pdrop=_a						,					embd_pdrop=_a						,					attn_pdrop=_a						,					layer_norm_epsilon=_a						,					initializer_range=_a						,					scale_attn_weights=_a						,					use_cache=_a						,					scale_attn_by_inverse_layer_idx=_a						,					reorder_and_upcast_attn=_a						,					)
							lowerCamelCase      =   GPTaLMHeadModel(_a					)





				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a = None						,					_a = None						,					):





							"""simple docstring"""
							lowerCamelCase      =   self.transformer.transformer.wte(_a					)
							lowerCamelCase      =   self.encode_prefix(_a					)
							lowerCamelCase      =   self.decode_prefix(_a					)
							lowerCamelCase      =   torch.cat((prefix_embeds, embedding_text)						,					dim=1					)

							if labels is not None:
										lowerCamelCase      =   self.get_dummy_token(input_ids.shape[0]						,					input_ids.device					)
										lowerCamelCase      =   torch.cat((dummy_token, input_ids)						,					dim=1					)
							lowerCamelCase      =   self.transformer(inputs_embeds=_a						,					labels=_a						,					attention_mask=_a					)
							if self.prefix_hidden_dim is not None:
										return out, hidden
							else:
										return out





				def   _lowerCAmelCase							(			self						,					_a						,					_a					):





							"""simple docstring"""
							return torch.zeros(_a						,					self.prefix_length						,					dtype=torch.intaa						,					device=_a					)





				def   _lowerCAmelCase							(			self						,					_a					):





							"""simple docstring"""
							return self.encode_prefix(_a					)





				@torch.no_grad()
				def   _lowerCAmelCase							(			self						,					_a						,					_a						,					_a					):





							"""simple docstring"""
							lowerCamelCase      =   torch.split(_a						,					1						,					dim=0					)
							lowerCamelCase      =   []
							lowerCamelCase      =   []
							for feature in features:
										lowerCamelCase      =   self.decode_prefix(feature.to(_a					)					)  # back to the clip feature
										# Only support beam search for now
										lowerCamelCase ,  lowerCamelCase      =   self.generate_beam(
										    input_embeds=_a						,					device=_a						,					eos_token_id=_a					)
										generated_tokens.append(output_tokens[0]					)
										generated_seq_lengths.append(seq_lengths[0]					)
							lowerCamelCase      =   torch.stack(_a					)
							lowerCamelCase      =   torch.stack(_a					)
							return generated_tokens, generated_seq_lengths





				@torch.no_grad()
				def   _lowerCAmelCase							(			self						,					_a=None						,					_a=None						,					_a=None						,					_a = 5						,					_a = 67						,					_a = 1.0						,					_a = None						,					):





							"""simple docstring"""
							lowerCamelCase      =   eos_token_id
							lowerCamelCase      =   None
							lowerCamelCase      =   None
							lowerCamelCase      =   torch.ones(_a						,					device=_a						,					dtype=torch.int					)
							lowerCamelCase      =   torch.zeros(_a						,					device=_a						,					dtype=torch.bool					)

							if input_embeds is not None:
										lowerCamelCase      =   input_embeds
							else:
										lowerCamelCase      =   self.transformer.transformer.wte(_a					)

							for i in range(_a					):
										lowerCamelCase      =   self.transformer(inputs_embeds=_a					)
										lowerCamelCase      =   outputs.logits
										lowerCamelCase      =   logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
										lowerCamelCase      =   logits.softmax(-1					).log()

										if scores is None:
													lowerCamelCase ,  lowerCamelCase      =   logits.topk(_a						,					-1					)
													lowerCamelCase      =   generated.expand(_a						,					*generated.shape[1:]					)
													lowerCamelCase ,  lowerCamelCase      =   next_tokens.permute(1						,					0					), scores.squeeze(0					)
													if tokens is None:
																lowerCamelCase      =   next_tokens
													else:
																lowerCamelCase      =   tokens.expand(_a						,					*tokens.shape[1:]					)
																lowerCamelCase      =   torch.cat((tokens, next_tokens)						,					dim=1					)
										else:
													lowerCamelCase      =   -float(np.inf					)
													lowerCamelCase      =   0
													lowerCamelCase      =   scores[:, None] + logits
													seq_lengths[~is_stopped] += 1
													lowerCamelCase      =   scores_sum / seq_lengths[:, None]
													lowerCamelCase ,  lowerCamelCase      =   scores_sum_average.view(-1					).topk(_a						,					-1					)
													lowerCamelCase      =   next_tokens // scores_sum.shape[1]
													lowerCamelCase      =   seq_lengths[next_tokens_source]
													lowerCamelCase      =   next_tokens % scores_sum.shape[1]
													lowerCamelCase      =   next_tokens.unsqueeze(1					)
													lowerCamelCase      =   tokens[next_tokens_source]
													lowerCamelCase      =   torch.cat((tokens, next_tokens)						,					dim=1					)
													lowerCamelCase      =   generated[next_tokens_source]
													lowerCamelCase      =   scores_sum_average * seq_lengths
													lowerCamelCase      =   is_stopped[next_tokens_source]

										lowerCamelCase      =   self.transformer.transformer.wte(next_tokens.squeeze()					).view(generated.shape[0]						,					1						,					-1					)
										lowerCamelCase      =   torch.cat((generated, next_token_embed)						,					dim=1					)
										lowerCamelCase      =   is_stopped + next_tokens.eq(_a					).squeeze()
										if is_stopped.all():
													break

							lowerCamelCase      =   scores / seq_lengths
							lowerCamelCase      =   scores.argsort(descending=_a					)
							# tokens tensors are already padded to max_seq_length
							lowerCamelCase      =   [tokens[i] for i in order]
							lowerCamelCase      =   torch.stack(_a						,					dim=0					)
							lowerCamelCase      =   torch.tensor([seq_lengths[i] for i in order]						,					dtype=seq_lengths.dtype					)
							return output_texts, seq_lengths



