





"""simple docstring"""





import csv
import json
import os

import datasets


# TODO: Add BibTeX citation
# Find for instance the citation on arxiv or on the dataset repo/website
lowerCAmelCase__        =      '''\
@InProceedings{huggingface:dataset,
title = {A great new dataset},
author={huggingface, Inc.
},
year={2020}
}
'''

# TODO: Add description of the dataset here
# You can copy an official description
lowerCAmelCase__        =      '''\
This new dataset is designed to solve this great NLP task and is crafted with a lot of care.
'''

# TODO: Add a link to an official homepage for the dataset here
lowerCAmelCase__        =      ''''''

# TODO: Add the licence for the dataset here if you can find it
lowerCAmelCase__        =      ''''''

# TODO: Add link to the official dataset URLs here
# The HuggingFace Datasets library doesn't host the datasets but only points to the original files.
# This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)
lowerCAmelCase__        =      {
    '''first_domain''': '''https://huggingface.co/great-new-dataset-first_domain.zip''',
    '''second_domain''': '''https://huggingface.co/great-new-dataset-second_domain.zip''',
}





class 							_lowerCamelCase   (							datasets.GeneratorBasedBuilder			):
						UpperCAmelCase_			=      datasets.Version("1.1.0"			)

						# This is an example of a dataset with multiple configurations.
						# If you don't want/need to define several sub-sets in your dataset,
						# just remove the BUILDER_CONFIG_CLASS and the BUILDER_CONFIGS attributes.

						# If you need to make complex sub-parts in the datasets with configurable options
						# You can create your own builder configuration class to store attribute, inheriting from datasets.BuilderConfig
						# BUILDER_CONFIG_CLASS = MyBuilderConfig

						# You will be able to load one or the other configurations in the following list with
						# data = datasets.load_dataset('my_dataset', 'first_domain')
						# data = datasets.load_dataset('my_dataset', 'second_domain')
						UpperCAmelCase_			=      [
						    datasets.BuilderConfig(name="first_domain"	,      version=_lowercase	,      description="This part of my dataset covers a first domain"			),
						    datasets.BuilderConfig(name="second_domain"	,      version=_lowercase	,      description="This part of my dataset covers a second domain"			),
						]

						UpperCAmelCase_			=      "first_domain"  # It's not mandatory to have a default configuration. Just use one if it make sense.


						def      snake_case_ (self		)			-> str:
								# TODO: This method specifies the datasets.DatasetInfo object which contains informations and typings for the dataset
								if self.config.name == "first_domain":  # This is the name of the configuration selected in BUILDER_CONFIGS above
										UpperCamelCase =						datasets.Features(
										    {
										        "sentence": datasets.Value("string"		),
										        "option1": datasets.Value("string"		),
										        "answer": datasets.Value("string"		)
										        # These are the features of your dataset like images, labels ...
										    }		)
								else:  # This is an example to show how to have different features for "first_domain" and "second_domain"
										UpperCamelCase =						datasets.Features(
										    {
										        "sentence": datasets.Value("string"		),
										        "option2": datasets.Value("string"		),
										        "second_domain_answer": datasets.Value("string"		)
										        # These are the features of your dataset like images, labels ...
										    }		)
								return datasets.DatasetInfo(
								    # This is the description that will appear on the datasets page.
								    description=_DESCRIPTION ,				features=__a ,				homepage=_HOMEPAGE ,				license=_LICENSE ,				citation=_CITATION ,				)


						def      snake_case_ (self ,				__a		)			-> Dict:
								# TODO: This method is tasked with downloading/extracting the data and defining the splits depending on the configuration
								# If several configurations are possible (listed in BUILDER_CONFIGS), the configuration selected by the user is in self.config.name

								# dl_manager is a datasets.download.DownloadManager that can be used to download and extract URLS
								# It can accept any type or nested list/dict and will give back the same structure with the url replaced with path to local files.
								# By default the archives will be extracted and a path to a cached folder where they are extracted is returned instead of the archive
								UpperCamelCase =						_URLS[self.config.name]
								UpperCamelCase =						dl_manager.download_and_extract(__a		)
								return [
								    datasets.SplitGenerator(
								        name=datasets.Split.TRAIN ,				gen_kwargs={
								            "filepath": os.path.join(__a ,				"train.jsonl"		),
								            "split": "train",
								        } ,				),
								    datasets.SplitGenerator(
								        name=datasets.Split.VALIDATION ,				gen_kwargs={
								            "filepath": os.path.join(__a ,				"dev.jsonl"		),
								            "split": "dev",
								        } ,				),
								    datasets.SplitGenerator(
								        name=datasets.Split.TEST ,				gen_kwargs={
								            "filepath": os.path.join(__a ,				"test.jsonl"		),
								            "split": "test"
								        } ,				),
								]




						def      snake_case_ (self ,				__a ,				__a		)			-> List[Any]:
								# TODO: This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.
								# The `key` is for legacy reasons (tfds) and is not important in itself, but must be unique for each example.
								with open(__a ,				encoding="utf-8"		) as f:
										for key, row in enumerate(__a		):
												UpperCamelCase =						json.loads(__a		)
												if self.config.name == "first_domain":
														# Yields examples as (key, example) tuples
														yield key, {
														    "sentence": data["sentence"],
														    "option1": data["option1"],
														    "answer": "" if split == "test" else data["answer"],
														}
												else:
														yield key, {
														    "sentence": data["sentence"],
														    "option2": data["option2"],
														    "second_domain_answer": "" if split == "test" else data["second_domain_answer"],
														}













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


lowerCAmelCase__        =      logging.get_logger(__name__)





class 							_lowerCamelCase   (							_lowercase			):


						def __init__(self ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = None ,				__a = (None, None) ,				__a = None ,				)			-> List[Any]:
								super().__init__(
								    model=__a ,				args=__a ,				data_collator=__a ,				train_dataset=__a ,				eval_dataset=__a ,				tokenizer=__a ,				model_init=__a ,				compute_metrics=__a ,				callbacks=__a ,				optimizers=__a ,				preprocess_logits_for_metrics=__a ,				)

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


						@staticmethod
						def      snake_case_ (__a		)			-> GenerationConfig:

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

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

								UpperCamelCase =						GenerationConfig.from_pretrained(__a ,				__a		)
								return gen_config


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

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


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

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


						def      snake_case_ (self ,				__a ,				__a ,				__a ,				__a = None ,				**__a ,				)			-> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:

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

								UpperCamelCase =						"labels" in inputs
								UpperCamelCase =						self._prepare_inputs(__a		)

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

								if len(__a		) == 0 and hasattr(self ,				"_gen_kwargs"		):
										UpperCamelCase =						self._gen_kwargs.copy()

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

								# 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:
										UpperCamelCase =						False

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

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

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

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

								return loss, generated_tokens, labels




						def      snake_case_ (self ,				__a ,				__a		)			-> 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
										UpperCamelCase =						(
										    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:
												UpperCamelCase =						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"		)

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






