





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


lowerCAmelCase_        =			logging.get_logger(__name__)



class 				_lowerCAmelCase			(	UpperCAmelCase_							):


  '''simple docstring'''







  def __init__(					self					:    List[str]					, UpperCamelCase					:    Union["PreTrainedModel", nn.Module] = None					, UpperCamelCase					:    "TrainingArguments" = None					, UpperCamelCase					:    Optional["DataCollator"] = None					, UpperCamelCase					:    Optional[Dataset] = None					, UpperCamelCase					:    Optional[Union[Dataset, Dict[str, Dataset]]] = None					, UpperCamelCase					:    Optional["PreTrainedTokenizerBase"] = None					, UpperCamelCase					:    Optional[Callable[[], "PreTrainedModel"]] = None					, UpperCamelCase					:    Optional[Callable[["EvalPrediction"], Dict]] = None					, UpperCamelCase					:    Optional[List["TrainerCallback"]] = None					, UpperCamelCase					:    Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)					, UpperCamelCase					:    Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None					, ):


    '''simple docstring'''





    super().__init__(
        model=UpperCamelCase					, args=UpperCamelCase					, data_collator=UpperCamelCase					, train_dataset=UpperCamelCase					, eval_dataset=UpperCamelCase					, tokenizer=UpperCamelCase					, model_init=UpperCamelCase					, compute_metrics=UpperCamelCase					, callbacks=UpperCamelCase					, optimizers=UpperCamelCase					, preprocess_logits_for_metrics=UpperCamelCase					, )

    # 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:
      _snake_case      : Dict														=		self.load_generation_config(self.args.generation_config							)
      _snake_case      : Any														=		gen_config







  @staticmethod
  def UpperCamelCase_       (					UpperCamelCase					:    Union[str, GenerationConfig]							):


    '''simple docstring'''





    if isinstance(UpperCamelCase					, UpperCamelCase							):
      return deepcopy(UpperCamelCase							)

    # str or Path
    _snake_case      : int														=		Path(UpperCamelCase							) if isinstance(UpperCamelCase					, UpperCamelCase							) else gen_config_arg
    _snake_case      : List[Any]														=		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():
      _snake_case      : int														=		pretrained_model_name.name
      _snake_case      : List[str]														=		pretrained_model_name.parent
    # dir path
    elif pretrained_model_name.is_dir():
      pass
    # model id or URL
    else:
      _snake_case      : Any														=		gen_config_arg

    _snake_case      : Optional[int]														=		GenerationConfig.from_pretrained(UpperCamelCase					, UpperCamelCase							)
    return gen_config







  def UpperCamelCase_       (					self					:    int					, UpperCamelCase					:    Optional[Dataset] = None					, UpperCamelCase					:    Optional[List[str]] = None					, UpperCamelCase					:    str = "eval"					, **UpperCamelCase					:    str					, ):


    '''simple docstring'''





    _snake_case      : Optional[int]														=		gen_kwargs.copy()
    if gen_kwargs.get('max_length'							) is None and gen_kwargs.get('max_new_tokens'							) is None:
      _snake_case      : Tuple														=		self.args.generation_max_length
    _snake_case      : str														=		(
        gen_kwargs['num_beams'] if gen_kwargs.get('num_beams'							) is not None else self.args.generation_num_beams
    )
    _snake_case      : int														=		gen_kwargs

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







  def UpperCamelCase_       (					self					:    Dict					, UpperCamelCase					:    Dataset					, UpperCamelCase					:    Optional[List[str]] = None					, UpperCamelCase					:    str = "test"					, **UpperCamelCase					:    Tuple					, ):


    '''simple docstring'''





    _snake_case      : List[Any]														=		gen_kwargs.copy()
    if gen_kwargs.get('max_length'							) is None and gen_kwargs.get('max_new_tokens'							) is None:
      _snake_case      : List[str]														=		self.args.generation_max_length
    _snake_case      : Optional[Any]														=		(
        gen_kwargs['num_beams'] if gen_kwargs.get('num_beams'							) is not None else self.args.generation_num_beams
    )
    _snake_case      : Union[str, Any]														=		gen_kwargs

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







  def UpperCamelCase_       (					self					:    Union[str, Any]					, UpperCamelCase					:    nn.Module					, UpperCamelCase					:    Dict[str, Union[torch.Tensor, Any]]					, UpperCamelCase					:    bool					, UpperCamelCase					:    Optional[List[str]] = None					, **UpperCamelCase					:    Optional[int]					, ):


    '''simple docstring'''





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

    _snake_case      : Dict														=		'labels' in inputs
    _snake_case      : Tuple														=		self._prepare_inputs(UpperCamelCase							)

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

    if len(UpperCamelCase							) == 0 and hasattr(self					, '_gen_kwargs'							):
      _snake_case      : Union[str, Any]														=		self._gen_kwargs.copy()

    if gen_kwargs.get('max_length'							) is None and gen_kwargs.get('max_new_tokens'							) is None:
      _snake_case      : Dict														=		self.model.config.max_length
    _snake_case      : int														=		(
        gen_kwargs['num_beams'] if gen_kwargs.get('num_beams'							) is not None else self.model.config.num_beams
    )
    _snake_case      : int														=		True if is_deepspeed_zeroa_enabled() else False
    _snake_case      : Optional[int]														=		(
        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
    ):
      _snake_case      : Union[str, Any]														=		{k: v for k, v in inputs.items() if k != 'decoder_input_ids'}
    _snake_case      : str														=		self.model.generate(**UpperCamelCase					, **UpperCamelCase							)

    # 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:
      _snake_case      : List[str]														=		False

    # Retrieves GenerationConfig from model.generation_config
    _snake_case      : List[Any]														=		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:
      _snake_case      : int														=		self._pad_tensors_to_max_len(UpperCamelCase					, gen_config.max_length							)
    elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
      _snake_case      : str														=		self._pad_tensors_to_max_len(UpperCamelCase					, gen_config.max_new_tokens + 1							)

    with torch.no_grad():
      if has_labels:
        with self.compute_loss_context_manager():
          _snake_case      : List[str]														=		model(**UpperCamelCase							)
        if self.label_smoother is not None:
          _snake_case      : Dict														=		self.label_smoother(UpperCamelCase					, inputs['labels']							).mean().detach()
        else:
          _snake_case      : str														=		(outputs['loss'] if isinstance(UpperCamelCase					, UpperCamelCase							) else outputs[0]).mean().detach()
      else:
        _snake_case      : List[Any]														=		None

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

    if has_labels:
      _snake_case      : Tuple														=		inputs['labels']
      if labels.shape[-1] < gen_config.max_length:
        _snake_case      : Dict														=		self._pad_tensors_to_max_len(UpperCamelCase					, gen_config.max_length							)
      elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
        _snake_case      : Optional[int]														=		self._pad_tensors_to_max_len(UpperCamelCase					, gen_config.max_new_tokens + 1							)
    else:
      _snake_case      : List[Any]														=		None

    return loss, generated_tokens, labels







  def UpperCamelCase_       (					self					:    Any					, UpperCamelCase					:    str					, UpperCamelCase					:    Tuple							):


    '''simple docstring'''





    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
      _snake_case      : Union[str, Any]														=		(
          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:
        _snake_case      : Dict														=		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'							)

    _snake_case      : List[str]														=		pad_token_id * torch.ones(
        (tensor.shape[0], max_length)					, dtype=tensor.dtype					, device=tensor.device							)
    _snake_case      : List[Any]														=		tensor
    return padded_tensor











import datasets

from .evaluate import (
    apply_no_ans_threshold,
    find_all_best_thresh,
    get_raw_scores,
    make_eval_dict,
    make_qid_to_has_ans,
    merge_eval,
)


lowerCAmelCase_        =			"""\
@inproceedings{Rajpurkar2016SQuAD10,
  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},
  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},
  booktitle={EMNLP},
  year={2016}
}
"""

lowerCAmelCase_        =			"""
This metric wrap the official scoring script for version 2 of the Stanford Question
Answering Dataset (SQuAD).

Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by
crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,
from the corresponding reading passage, or the question might be unanswerable.

SQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions
written adversarially by crowdworkers to look similar to answerable ones.
To do well on SQuAD2.0, systems must not only answer questions when possible, but also
determine when no answer is supported by the paragraph and abstain from answering.
"""

lowerCAmelCase_        =			"""
Computes SQuAD v2 scores (F1 and EM).
Args:
    predictions: List of triple for question-answers to score with the following elements:
        - the question-answer 'id' field as given in the references (see below)
        - the text of the answer
        - the probability that the question has no answer
    references: List of question-answers dictionaries with the following key-values:
            - 'id': id of the question-answer pair (see above),
            - 'answers': a list of Dict {'text': text of the answer as a string}
    no_answer_threshold: float
        Probability threshold to decide that a question has no answer.
Returns:
    'exact': Exact match (the normalized answer exactly match the gold answer)
    'f1': The F-score of predicted tokens versus the gold answer
    'total': Number of score considered
    'HasAns_exact': Exact match (the normalized answer exactly match the gold answer)
    'HasAns_f1': The F-score of predicted tokens versus the gold answer
    'HasAns_total': Number of score considered
    'NoAns_exact': Exact match (the normalized answer exactly match the gold answer)
    'NoAns_f1': The F-score of predicted tokens versus the gold answer
    'NoAns_total': Number of score considered
    'best_exact': Best exact match (with varying threshold)
    'best_exact_thresh': No-answer probability threshold associated to the best exact match
    'best_f1': Best F1 (with varying threshold)
    'best_f1_thresh': No-answer probability threshold associated to the best F1
Examples:

    >>> predictions = [{'prediction_text': '1976', 'id': '56e10a3be3433e1400422b22', 'no_answer_probability': 0.}]
    >>> references = [{'answers': {'answer_start': [97], 'text': ['1976']}, 'id': '56e10a3be3433e1400422b22'}]
    >>> squad_v2_metric = datasets.load_metric(\"squad_v2\")
    >>> results = squad_v2_metric.compute(predictions=predictions, references=references)
    >>> print(results)
    {'exact': 100.0, 'f1': 100.0, 'total': 1, 'HasAns_exact': 100.0, 'HasAns_f1': 100.0, 'HasAns_total': 1, 'best_exact': 100.0, 'best_exact_thresh': 0.0, 'best_f1': 100.0, 'best_f1_thresh': 0.0}
"""



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,			_KWARGS_DESCRIPTION							)
class 				_lowerCAmelCase			(	datasets.Metric							):


  '''simple docstring'''







  def UpperCamelCase_       (					self					:    str							):


    '''simple docstring'''





    return datasets.MetricInfo(
        description=_DESCRIPTION					, citation=_CITATION					, inputs_description=_KWARGS_DESCRIPTION					, features=datasets.Features(
            {
                'predictions': {
                    'id': datasets.Value('string'							),
                    'prediction_text': datasets.Value('string'							),
                    'no_answer_probability': datasets.Value('float32'							),
                },
                'references': {
                    'id': datasets.Value('string'							),
                    'answers': datasets.features.Sequence(
                        {'text': datasets.Value('string'							), 'answer_start': datasets.Value('int32'							)}							),
                },
            }							)					, codebase_urls=['https://rajpurkar.github.io/SQuAD-explorer/']					, reference_urls=['https://rajpurkar.github.io/SQuAD-explorer/']					, )







  def UpperCamelCase_       (					self					:    Optional[Any]					, UpperCamelCase					:    List[str]					, UpperCamelCase					:    Optional[Any]					, UpperCamelCase					:    str=1.0							):


    '''simple docstring'''





    _snake_case      : Dict														=		{p['id']: p['no_answer_probability'] for p in predictions}
    _snake_case      : Dict														=		[{'paragraphs': [{'qas': references}]}]
    _snake_case      : Dict														=		{p['id']: p['prediction_text'] for p in predictions}

    _snake_case      : Optional[Any]														=		make_qid_to_has_ans(UpperCamelCase							)  # maps qid to True/False
    _snake_case      : str														=		[k for k, v in qid_to_has_ans.items() if v]
    _snake_case      : Optional[Any]														=		[k for k, v in qid_to_has_ans.items() if not v]

    _snake_case							,    _snake_case      : Union[str, Any]														=		get_raw_scores(UpperCamelCase					, UpperCamelCase							)
    _snake_case      : int														=		apply_no_ans_threshold(UpperCamelCase					, UpperCamelCase					, UpperCamelCase					, UpperCamelCase							)
    _snake_case      : List[Any]														=		apply_no_ans_threshold(UpperCamelCase					, UpperCamelCase					, UpperCamelCase					, UpperCamelCase							)
    _snake_case      : Optional[Any]														=		make_eval_dict(UpperCamelCase					, UpperCamelCase							)

    if has_ans_qids:
      _snake_case      : Optional[int]														=		make_eval_dict(UpperCamelCase					, UpperCamelCase					, qid_list=UpperCamelCase							)
      merge_eval(UpperCamelCase					, UpperCamelCase					, 'HasAns'							)
    if no_ans_qids:
      _snake_case      : List[str]														=		make_eval_dict(UpperCamelCase					, UpperCamelCase					, qid_list=UpperCamelCase							)
      merge_eval(UpperCamelCase					, UpperCamelCase					, 'NoAns'							)
    find_all_best_thresh(UpperCamelCase					, UpperCamelCase					, UpperCamelCase					, UpperCamelCase					, UpperCamelCase					, UpperCamelCase							)
    return dict(UpperCamelCase							)




