





"""simple docstring"""


import textwrap

import pyarrow as pa
import pytest

from datasets import Features, Image
from datasets.packaged_modules.text.text import Text

from ..utils import require_pil







@pytest.fixture
def       UpperCAmelCase					(	UpperCAmelCase      )       ->   int:
     snake_case_       =      tmp_path / 'text.txt'
     snake_case_       =      textwrap.dedent(
         '\\n        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.\n        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.\n        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.\n        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.\n\n        Second paragraph:\n        Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua.\n        Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat.\n        Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur.\n        Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum.\n        '      )
     with open(UpperCAmelCase		,      'w'		,      encoding='utf-8'      ) as f:
          f.write(UpperCAmelCase      )
     return str(UpperCAmelCase      )







@pytest.fixture
def       UpperCAmelCase					(	UpperCAmelCase		,      UpperCAmelCase      )       ->   Union[str, Any]:
     snake_case_       =      tmp_path / 'text_with_image.txt'
     with open(UpperCAmelCase		,      'w'		,      encoding='utf-8'      ) as f:
          f.write(UpperCAmelCase      )
     return str(UpperCAmelCase      )







@pytest.mark.parametrize('keep_linebreaks'		,      [True, False]      )
def       UpperCAmelCase					(	UpperCAmelCase		,      UpperCAmelCase      )       ->   Tuple:
     with open(UpperCAmelCase		,      encoding='utf-8'      ) as f:
          snake_case_       =      f.read().splitlines(keepends=UpperCAmelCase      )
     snake_case_       =      Text(keep_linebreaks=UpperCAmelCase		,      encoding='utf-8'      )
     snake_case_       =      text._generate_tables([[text_file]]      )
     snake_case_       =      pa.concat_tables([table for _, table in generator]      ).to_pydict()['text']
     assert generated_content == expected_content







@require_pil
def       UpperCAmelCase					(	UpperCAmelCase      )       ->   Any:
     with open(UpperCAmelCase		,      encoding='utf-8'      ) as f:
          snake_case_       =      f.read().splitlines()[0]
     snake_case_       =      Text(encoding='utf-8'		,      features=Features({'image': Image()}      )      )
     snake_case_       =      text._generate_tables([[text_file_with_image]]      )
     snake_case_       =      pa.concat_tables([table for _, table in generator]      )
     assert pa_table.schema.field('image'      ).type == Image()()
     snake_case_       =      pa_table.to_pydict()['image']
     assert generated_content == [{"path": image_file, "bytes": None}]







@pytest.mark.parametrize('sample_by'		,      ['line', 'paragraph', 'document']      )
def       UpperCAmelCase					(	UpperCAmelCase		,      UpperCAmelCase      )       ->   Dict:
     with open(UpperCAmelCase		,      encoding='utf-8'      ) as f:
          snake_case_       =      f.read()
     if sample_by == "line":
          snake_case_       =      expected_content.splitlines()
     elif sample_by == "paragraph":
          snake_case_       =      expected_content.split('\n\n'      )
     elif sample_by == "document":
          snake_case_       =      [expected_content]
     snake_case_       =      Text(sample_by=UpperCAmelCase		,      encoding='utf-8'		,      chunksize=100      )
     snake_case_       =      text._generate_tables([[text_file]]      )
     snake_case_       =      pa.concat_tables([table for _, table in generator]      ).to_pydict()['text']
     assert generated_content == expected_content


"""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


__UpperCamelCase								=    logging.get_logger(__name__)




class 				UpperCamelCase      (       lowerCAmelCase__     ):







       def __init__(				self,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = (None, None),							lowerCAmelCase__ = None,							)   ->   Dict:
            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:
                 snake_case_       =      self.load_generation_config(self.args.generation_config)
                 snake_case_       =      gen_config







       @staticmethod
       def 	a_		(				lowerCAmelCase__)   ->   GenerationConfig:

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

            # str or Path
            snake_case_       =      Path(lowerCAmelCase__) if isinstance(lowerCAmelCase__,							lowerCAmelCase__) else gen_config_arg
            snake_case_       =      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_       =      pretrained_model_name.name
                 snake_case_       =      pretrained_model_name.parent
            # dir path
            elif pretrained_model_name.is_dir():
                 pass
            # model id or URL
            else:
                 snake_case_       =      gen_config_arg

            snake_case_       =      GenerationConfig.from_pretrained(lowerCAmelCase__,							lowerCAmelCase__)
            return gen_config







       def 	a_		(				self,							lowerCAmelCase__ = None,							lowerCAmelCase__ = None,							lowerCAmelCase__ = "eval",							**lowerCAmelCase__,							)   ->   Dict[str, float]:
            snake_case_       =      gen_kwargs.copy()
            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.args.generation_max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
            )
            snake_case_       =      gen_kwargs

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







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__ = None,							lowerCAmelCase__ = "test",							**lowerCAmelCase__,							)   ->   "PredictionOutput":
            snake_case_       =      gen_kwargs.copy()
            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.args.generation_max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.args.generation_num_beams
            )
            snake_case_       =      gen_kwargs

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







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__,							lowerCAmelCase__ = None,							**lowerCAmelCase__,							)   ->   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__)

            snake_case_       =      'labels' in inputs
            snake_case_       =      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'):
                 snake_case_       =      self._gen_kwargs.copy()

            if gen_kwargs.get('max_length') is None and gen_kwargs.get('max_new_tokens') is None:
                 snake_case_       =      self.model.config.max_length
            snake_case_       =      (
                gen_kwargs['num_beams'] if gen_kwargs.get('num_beams') is not None else self.model.config.num_beams
            )
            snake_case_       =      True if is_deepspeed_zeroa_enabled() else False
            snake_case_       =      (
                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_       =      {k: v for k, v in inputs.items() if k != 'decoder_input_ids'}
            snake_case_       =      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:
                 snake_case_       =      False

            # Retrieves GenerationConfig from model.generation_config
            snake_case_       =      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_       =      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:
                 snake_case_       =      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():
                           snake_case_       =      model(**lowerCAmelCase__)
                      if self.label_smoother is not None:
                           snake_case_       =      self.label_smoother(lowerCAmelCase__,							inputs['labels']).mean().detach()
                      else:
                           snake_case_       =      (outputs['loss'] if isinstance(lowerCAmelCase__,							lowerCAmelCase__) else outputs[0]).mean().detach()
                 else:
                      snake_case_       =      None

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

            if has_labels:
                 snake_case_       =      inputs['labels']
                 if labels.shape[-1] < gen_config.max_length:
                      snake_case_       =      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:
                      snake_case_       =      self._pad_tensors_to_max_len(lowerCAmelCase__,							gen_config.max_new_tokens + 1)
            else:
                 snake_case_       =      None

            return loss, generated_tokens, labels







       def 	a_		(				self,							lowerCAmelCase__,							lowerCAmelCase__)   ->   Any:
            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_       =      (
                     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_       =      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_       =      pad_token_id * torch.ones(
                (tensor.shape[0], max_length),							dtype=tensor.dtype,							device=tensor.device)
            snake_case_       =      tensor
            return padded_tensor

