






"""simple docstring"""

from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import Dataset

from .deepspeed import is_deepspeed_zeroa_enabled
from .generation.configuration_utils import GenerationConfig
from .trainer import Trainer
from .utils import logging


if TYPE_CHECKING:
			from .data.data_collator import DataCollator
			from .modeling_utils import PreTrainedModel
			from .tokenization_utils_base import PreTrainedTokenizerBase
			from .trainer_callback import TrainerCallback
			from .trainer_utils import EvalPrediction, PredictionOutput
			from .training_args import TrainingArguments


lowercase           =							logging.get_logger(__name__)

class     A_				(   snake_case_	):







				def __init__(				self				:    Optional[Any]    ,     __lowerCamelCase				:    Union["PreTrainedModel", nn.Module] = None    ,     __lowerCamelCase				:    "TrainingArguments" = None    ,     __lowerCamelCase				:    Optional["DataCollator"] = None    ,     __lowerCamelCase				:    Optional[Dataset] = None    ,     __lowerCamelCase				:    Optional[Union[Dataset, Dict[str, Dataset]]] = None    ,     __lowerCamelCase				:    Optional["PreTrainedTokenizerBase"] = None    ,     __lowerCamelCase				:    Optional[Callable[[], "PreTrainedModel"]] = None    ,     __lowerCamelCase				:    Optional[Callable[["EvalPrediction"], Dict]] = None    ,     __lowerCamelCase				:    Optional[List["TrainerCallback"]] = None    ,     __lowerCamelCase				:    Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)    ,     __lowerCamelCase				:    Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None    ,     )					->       Tuple:
										super().__init__(
										    model=__lowerCamelCase    ,     args=__lowerCamelCase    ,     data_collator=__lowerCamelCase    ,     train_dataset=__lowerCamelCase    ,     eval_dataset=__lowerCamelCase    ,     tokenizer=__lowerCamelCase    ,     model_init=__lowerCamelCase    ,     compute_metrics=__lowerCamelCase    ,     callbacks=__lowerCamelCase    ,     optimizers=__lowerCamelCase    ,     preprocess_logits_for_metrics=__lowerCamelCase    ,     )

										# Override self.model.generation_config if a GenerationConfig is specified in args.
										# Priority: args.generation_config > model.generation_config > default GenerationConfig.
										if self.args.generation_config is not None:
																__magic_name__      =  self.load_generation_config(self.args.generation_config						)
																__magic_name__      =  gen_config







				@staticmethod
				def      _snake_case       (				__lowerCamelCase				:    Union[str, GenerationConfig]						)					->       GenerationConfig:

										# GenerationConfig provided, nothing to do
										if isinstance(__lowerCamelCase    ,     __lowerCamelCase						):
																return deepcopy(__lowerCamelCase						)

										# str or Path
										__magic_name__      =  Path(__lowerCamelCase						) if isinstance(__lowerCamelCase    ,     __lowerCamelCase						) else gen_config_arg
										__magic_name__      =  None

										# Figuring if it is path pointing to a file, pointing to a directory or else a model id or URL
										# This step is required in order to determine config_file_name
										if pretrained_model_name.is_file():
																__magic_name__      =  pretrained_model_name.name
																__magic_name__      =  pretrained_model_name.parent
										# dir path
										elif pretrained_model_name.is_dir():
																pass
										# model id or URL
										else:
																__magic_name__      =  gen_config_arg

										__magic_name__      =  GenerationConfig.from_pretrained(__lowerCamelCase    ,     __lowerCamelCase						)
										return gen_config







				def      _snake_case       (				self				:    Any    ,     __lowerCamelCase				:    Optional[Dataset] = None    ,     __lowerCamelCase				:    Optional[List[str]] = None    ,     __lowerCamelCase				:    str = "eval"    ,     **__lowerCamelCase				:    Union[str, Any]    ,     )					->       Dict[str, float]:
										__magic_name__      =  gen_kwargs.copy()
										if gen_kwargs.get("max_length"						) is None and gen_kwargs.get("max_new_tokens"						) is None:
																__magic_name__      =  self.args.generation_max_length
										__magic_name__      =  (
										    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"						) is not None else self.args.generation_num_beams
										)
										__magic_name__      =  gen_kwargs

										return super().evaluate(__lowerCamelCase    ,     ignore_keys=__lowerCamelCase    ,     metric_key_prefix=__lowerCamelCase						)







				def      _snake_case       (				self				:    Any    ,     __lowerCamelCase				:    Dataset    ,     __lowerCamelCase				:    Optional[List[str]] = None    ,     __lowerCamelCase				:    str = "test"    ,     **__lowerCamelCase				:    List[Any]    ,     )					->       "PredictionOutput":
										__magic_name__      =  gen_kwargs.copy()
										if gen_kwargs.get("max_length"						) is None and gen_kwargs.get("max_new_tokens"						) is None:
																__magic_name__      =  self.args.generation_max_length
										__magic_name__      =  (
										    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"						) is not None else self.args.generation_num_beams
										)
										__magic_name__      =  gen_kwargs

										return super().predict(__lowerCamelCase    ,     ignore_keys=__lowerCamelCase    ,     metric_key_prefix=__lowerCamelCase						)







				def      _snake_case       (				self				:    Optional[Any]    ,     __lowerCamelCase				:    nn.Module    ,     __lowerCamelCase				:    Dict[str, Union[torch.Tensor, Any]]    ,     __lowerCamelCase				:    bool    ,     __lowerCamelCase				:    Optional[List[str]] = None    ,     **__lowerCamelCase				:    int    ,     )					->       Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:

										if not self.args.predict_with_generate or prediction_loss_only:
																return super().prediction_step(
																    __lowerCamelCase    ,     __lowerCamelCase    ,     prediction_loss_only=__lowerCamelCase    ,     ignore_keys=__lowerCamelCase						)

										__magic_name__      =  "labels" in inputs
										__magic_name__      =  self._prepare_inputs(__lowerCamelCase						)

										# XXX: adapt synced_gpus for fairscale as well
										# Priority (handled in generate):
										# gen_kwargs > model.generation_config > default GenerationConfig()

										if len(__lowerCamelCase						) == 0 and hasattr(self    ,     "_gen_kwargs"						):
																__magic_name__      =  self._gen_kwargs.copy()

										if gen_kwargs.get("max_length"						) is None and gen_kwargs.get("max_new_tokens"						) is None:
																__magic_name__      =  self.model.config.max_length
										__magic_name__      =  (
										    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"						) is not None else self.model.config.num_beams
										)
										__magic_name__      =  True if is_deepspeed_zeroa_enabled() else False
										__magic_name__      =  (
										    gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus"						) is not None else default_synced_gpus
										)

										# If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate
										# (otherwise, it would continue generating from the padded `decoder_input_ids`)
										if (
										    "labels" in inputs
										    and "decoder_input_ids" in inputs
										    and inputs["labels"].shape == inputs["decoder_input_ids"].shape
										):
																__magic_name__      =  {k: v for k, v in inputs.items() if k != "decoder_input_ids"}
										__magic_name__      =  self.model.generate(**__lowerCamelCase    ,     **__lowerCamelCase						)

										# Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
										# TODO: remove this hack when the legacy code that initializes generation_config from a model config is
										# removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183
										if self.model.generation_config._from_model_config:
																__magic_name__      =  False

										# Retrieves GenerationConfig from model.generation_config
										__magic_name__      =  self.model.generation_config
										# in case the batch is shorter than max length, the output should be padded
										if generated_tokens.shape[-1] < gen_config.max_length:
																__magic_name__      =  self._pad_tensors_to_max_len(__lowerCamelCase    ,     gen_config.max_length						)
										elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
																__magic_name__      =  self._pad_tensors_to_max_len(__lowerCamelCase    ,     gen_config.max_new_tokens + 1						)

										with torch.no_grad():
																if has_labels:
																						with self.compute_loss_context_manager():
																												__magic_name__      =  model(**__lowerCamelCase						)
																						if self.label_smoother is not None:
																												__magic_name__      =  self.label_smoother(__lowerCamelCase    ,     inputs["labels"]						).mean().detach()
																						else:
																												__magic_name__      =  (outputs["loss"] if isinstance(__lowerCamelCase    ,     __lowerCamelCase						) else outputs[0]).mean().detach()
																else:
																						__magic_name__      =  None

										if self.args.prediction_loss_only:
																return loss, None, None

										if has_labels:
																__magic_name__      =  inputs["labels"]
																if labels.shape[-1] < gen_config.max_length:
																						__magic_name__      =  self._pad_tensors_to_max_len(__lowerCamelCase    ,     gen_config.max_length						)
																elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
																						__magic_name__      =  self._pad_tensors_to_max_len(__lowerCamelCase    ,     gen_config.max_new_tokens + 1						)
										else:
																__magic_name__      =  None

										return loss, generated_tokens, labels







				def      _snake_case       (				self				:    Optional[Any]    ,     __lowerCamelCase				:    Dict    ,     __lowerCamelCase				:    Union[str, Any]						)					->       str:
										if self.tokenizer is not None and hasattr(self.tokenizer    ,     "pad_token_id"						):
																# If PAD token is not defined at least EOS token has to be defined
																__magic_name__      =  (
																    self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
																)
										else:
																if self.model.config.pad_token_id is not None:
																						__magic_name__      =  self.model.config.pad_token_id
																else:
																						raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors"						)

										__magic_name__      =  pad_token_id * torch.ones(
										    (tensor.shape[0], max_length)    ,     dtype=tensor.dtype    ,     device=tensor.device						)
										__magic_name__      =  tensor
										return padded_tensor











"""simple docstring"""

import os
import tempfile
import unittest
from typing import List

from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
from transformers.testing_utils import require_tokenizers
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.tokenization_utils_base import PreTrainedTokenizerBase

@require_tokenizers
class     A_				(   unittest.TestCase	):
				UpperCAmelCase__										=					EsmTokenizer







				def      _snake_case       (				self				:    List[str]						)					->       Dict:
										super().setUp()
										__magic_name__      =  tempfile.mkdtemp()
										# fmt: off
										__magic_name__      =  ["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"]  # noqa: E501
										# fmt: on
										__magic_name__      =  os.path.join(self.tmpdirname    ,     VOCAB_FILES_NAMES["vocab_file"]						)
										with open(self.vocab_file    ,     "w"    ,     encoding="utf-8"						) as vocab_writer:
																vocab_writer.write("".join([x + "\n" for x in vocab_tokens]						)						)







				def      _snake_case       (				self				:    Optional[Any]    ,     **__lowerCamelCase				:    str						)					->       List[PreTrainedTokenizerBase]:
										return [self.get_tokenizer(**__lowerCamelCase						)]







				def      _snake_case       (				self				:    Dict    ,     **__lowerCamelCase				:    int						)					->       PreTrainedTokenizer:
										return self.tokenizer_class.from_pretrained(self.tmpdirname    ,     **__lowerCamelCase						)







				def      _snake_case       (				self				:    str						)					->       Optional[int]:
										__magic_name__      =  self.tokenizer_class(self.vocab_file						)

										__magic_name__      =  tokenizer.tokenize("LAGVS"						)
										self.assertListEqual(__lowerCamelCase    ,     ["L", "A", "G", "V", "S"]						)
										self.assertListEqual(tokenizer.convert_tokens_to_ids(__lowerCamelCase						)    ,     [4, 5, 6, 7, 8]						)







				def      _snake_case       (				self				:    int						)					->       Dict:
										__magic_name__      =  self.tokenizer_class(self.vocab_file						)

										__magic_name__      =  "LAGVS"
										self.assertListEqual(tokenizer.encode(__lowerCamelCase						)    ,     [0, 4, 5, 6, 7, 8, 2]						)







				def      _snake_case       (				self				:    List[Any]						)					->       Tuple:
										__magic_name__      =  self.tokenizer_class(self.vocab_file						)

										__magic_name__      =  ["LAGVS", "WCB"]
										__magic_name__      =  tokenizer(__lowerCamelCase    ,     padding=__lowerCamelCase						)["input_ids"]

										self.assertListEqual(__lowerCamelCase    ,     [[0, 4, 5, 6, 7, 8, 2], [0, 2_2, 2_3, 2_5, 2]]						)







				def      _snake_case       (				self				:    Union[str, Any]						)					->       List[str]:
										__magic_name__      =  self.tokenizer_class(self.vocab_file						)

										__magic_name__      =  ["LAGVS", "WCB"]
										__magic_name__      =  tokenizer(__lowerCamelCase    ,     padding=__lowerCamelCase						)["input_ids"]

										self.assertListEqual(__lowerCamelCase    ,     [[0, 4, 5, 6, 7, 8, 2], [0, 2_2, 2_3, 2_5, 2, 1, 1]]						)







				def      _snake_case       (				self				:    Union[str, Any]						)					->       Dict:
										__magic_name__      =  self.get_tokenizers(fast=__lowerCamelCase						)
										for tokenizer in tokenizers:
																with self.subTest(f'''{tokenizer.__class__.__name__}'''						):
																						__magic_name__      =  "<unk>"
																						__magic_name__      =  "<mask>"

																						__magic_name__      =  tokenizer.tokenize(__lowerCamelCase						)
																						__magic_name__      =  tokenizer.tokenize(__lowerCamelCase						)

																						self.assertEqual(len(__lowerCamelCase						)    ,     1						)
																						self.assertEqual(len(__lowerCamelCase						)    ,     1						)
																						self.assertEqual(token_a[0]    ,     __lowerCamelCase						)
																						self.assertEqual(token_a[0]    ,     __lowerCamelCase						)



