

"""simple docstring"""


import itertools
from dataclasses import dataclass
from typing import Optional

import pandas as pd
import pyarrow as pa

import datasets
from datasets.table import table_cast





@dataclass
class        snake_case      (							datasets.BuilderConfig    ):
				SCREAMING_SNAKE_CASE_  : Dict   =    None





class        snake_case      (							datasets.ArrowBasedBuilder    ):
				SCREAMING_SNAKE_CASE_  : Optional[int]   =    PandasConfig

				def 				lowercase_	(      self				:							Any)->  Tuple:



						'''simple docstring'''


						return datasets.DatasetInfo(features=self.config.features)

				def 				lowercase_	(      self				:							str	,    UpperCamelCase__				:							List[str])->  Optional[Any]:



						'''simple docstring'''


						if not self.config.data_files:
								raise ValueError(f"At least one data file must be specified, but got data_files={self.config.data_files}")
						__lowerCAmelCase:			int        =     dl_manager.download_and_extract(self.config.data_files)
						if isinstance(__lowercase	,    (str, list, tuple)):
								__lowerCAmelCase:			Dict        =     data_files
								if isinstance(__lowercase	,    __lowercase):
										__lowerCAmelCase:			Optional[Any]        =     [files]
								# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
								__lowerCAmelCase:			int        =     [dl_manager.iter_files(__lowercase) for file in files]
								return [datasets.SplitGenerator(name=datasets.Split.TRAIN	,    gen_kwargs={"files": files})]
						__lowerCAmelCase:			str        =     []
						for split_name, files in data_files.items():
								if isinstance(__lowercase	,    __lowercase):
										__lowerCAmelCase:			Union[str, Any]        =     [files]
								# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
								__lowerCAmelCase:			Any        =     [dl_manager.iter_files(__lowercase) for file in files]
								splits.append(datasets.SplitGenerator(name=__lowercase	,    gen_kwargs={"files": files}))
						return splits

				def 				lowercase_	(      self				:							List[str]	,    UpperCamelCase__				:							pa.Table)->  List[str]:



						'''simple docstring'''


						if self.config.features is not None:
								# more expensive cast to support nested features with keys in a different order
								# allows str <-> int/float or str to Audio for example
								__lowerCAmelCase:			List[str]        =     table_cast(__lowercase	,    self.config.features.arrow_schema)
						return pa_table




				def 				lowercase_	(      self				:							Any	,    UpperCamelCase__				:							Tuple)->  Dict:



						'''simple docstring'''


						for i, file in enumerate(itertools.chain.from_iterable(__lowercase)):
								with open(__lowercase	,    "rb") as f:
										__lowerCAmelCase:			str        =     pa.Table.from_pandas(pd.read_pickle(__lowercase))
										yield i, self._cast_table(__lowercase)





'''simple docstring'''

import enum
import warnings

from ..tokenization_utils import TruncationStrategy
from ..utils import add_end_docstrings, is_tf_available, is_torch_available, logging
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_tf_available():
       import tensorflow as tf

       from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

if is_torch_available():
       from ..models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

UpperCAmelCase					 =	logging.get_logger(__name__)


class 	lowerCAmelCase  (   enum.Enum ):
       lowerCAmelCase_												=			0
       lowerCAmelCase_												=			1






@add_end_docstrings(A )
class 	lowerCAmelCase  (   A ):
       lowerCAmelCase_												=			"generated"







       def __init__(				self      :			Tuple    ,   *__lowercase      :			Tuple    ,   **__lowercase      :			Optional[int]   ):







            """simple docstring"""







            super().__init__(*__lowercase    ,   **__lowercase   )

            self.check_model_type(
                TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
                if self.framework == 'tf'
                else MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING   )







       def       snake_case (				self      :			Dict    ,   __lowercase      :			Any=None    ,   __lowercase      :			int=None    ,   __lowercase      :			List[str]=None    ,   __lowercase      :			int=None    ,   __lowercase      :			str=None    ,   __lowercase      :			List[Any]=None    ,   **__lowercase      :			Union[str, Any]    ,   ):







            """simple docstring"""







            __lowercase      		={}
            if truncation is not None:
                 __lowercase      		=truncation

            __lowercase      		=generate_kwargs

            __lowercase      		={}
            if return_tensors is not None and return_type is None:
                 __lowercase      		=ReturnType.TENSORS if return_tensors else ReturnType.TEXT
            if return_type is not None:
                 __lowercase      		=return_type

            if clean_up_tokenization_spaces is not None:
                 __lowercase      		=clean_up_tokenization_spaces

            if stop_sequence is not None:
                 __lowercase      		=self.tokenizer.encode(__lowercase    ,   add_special_tokens=__lowercase   )
                 if len(__lowercase   ) > 1:
                      warnings.warn(
                          'Stopping on a multiple token sequence is not yet supported on transformers. The first token of'
                          ' the stop sequence will be used as the stop sequence string in the interim.'   )
                 __lowercase      		=stop_sequence_ids[0]

            return preprocess_params, forward_params, postprocess_params







       def       snake_case (				self      :			List[Any]    ,   __lowercase      :			int    ,   __lowercase      :			int    ,   __lowercase      :			int   ):







            """simple docstring"""







            return True







       def       snake_case (				self      :			Optional[Any]    ,   *__lowercase      :			Optional[int]    ,   __lowercase      :			Union[str, Any]   ):







            """simple docstring"""







            __lowercase      		=self.model.config.prefix if self.model.config.prefix is not None else ''
            if isinstance(args[0]    ,   __lowercase   ):
                 if self.tokenizer.pad_token_id is None:
                      raise ValueError('Please make sure that the tokenizer has a pad_token_id when using a batch input'   )
                 __lowercase      		=([prefix + arg for arg in args[0]],)
                 __lowercase      		=True

            elif isinstance(args[0]    ,   __lowercase   ):
                 __lowercase      		=(prefix + args[0],)
                 __lowercase      		=False
            else:
                 raise ValueError(
                     f''' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`'''   )
            __lowercase      		=self.tokenizer(*__lowercase    ,   padding=__lowercase    ,   truncation=__lowercase    ,   return_tensors=self.framework   )
            # This is produced by tokenizers but is an invalid generate kwargs
            if "token_type_ids" in inputs:
                 del inputs["token_type_ids"]
            return inputs







       def __call__(				self      :			Any    ,   *__lowercase      :			int    ,   **__lowercase      :			Optional[Any]   ):







            """simple docstring"""







            __lowercase      		=super().__call__(*__lowercase    ,   **__lowercase   )
            if (
                isinstance(args[0]    ,   __lowercase   )
                and all(isinstance(__lowercase    ,   __lowercase   ) for el in args[0]   )
                and all(len(__lowercase   ) == 1 for res in result   )
            ):
                 return [res[0] for res in result]
            return result







       def       snake_case (				self      :			List[str]    ,   __lowercase      :			Tuple    ,   __lowercase      :			List[Any]=TruncationStrategy.DO_NOT_TRUNCATE    ,   **__lowercase      :			Dict   ):







            """simple docstring"""







            __lowercase      		=self._parse_and_tokenize(__lowercase    ,   truncation=__lowercase    ,   **__lowercase   )
            return inputs







       def       snake_case (				self      :			Optional[int]    ,   __lowercase      :			Tuple    ,   **__lowercase      :			Tuple   ):







            """simple docstring"""







            if self.framework == "pt":
                 __lowercase		,  __lowercase      		=model_inputs['input_ids'].shape
            elif self.framework == "tf":
                 __lowercase		,  __lowercase      		=tf.shape(model_inputs['input_ids']   ).numpy()

            __lowercase      		=generate_kwargs.get('min_length'    ,   self.model.config.min_length   )
            __lowercase      		=generate_kwargs.get('max_length'    ,   self.model.config.max_length   )
            self.check_inputs(__lowercase    ,   generate_kwargs['min_length']    ,   generate_kwargs['max_length']   )
            __lowercase      		=self.model.generate(**__lowercase    ,   **__lowercase   )
            __lowercase      		=output_ids.shape[0]
            if self.framework == "pt":
                 __lowercase      		=output_ids.reshape(__lowercase    ,   out_b // in_b    ,   *output_ids.shape[1:]   )
            elif self.framework == "tf":
                 __lowercase      		=tf.reshape(__lowercase    ,   (in_b, out_b // in_b, *output_ids.shape[1:])   )
            return {"output_ids": output_ids}







       def       snake_case (				self      :			str    ,   __lowercase      :			Optional[Any]    ,   __lowercase      :			Optional[int]=ReturnType.TEXT    ,   __lowercase      :			Union[str, Any]=False   ):







            """simple docstring"""







            __lowercase      		=[]
            for output_ids in model_outputs["output_ids"][0]:
                 if return_type == ReturnType.TENSORS:
                      __lowercase      		={f'''{self.return_name}_token_ids''': output_ids}
                 elif return_type == ReturnType.TEXT:
                      __lowercase      		={
                          f'''{self.return_name}_text''': self.tokenizer.decode(
                              __lowercase    ,   skip_special_tokens=__lowercase    ,   clean_up_tokenization_spaces=__lowercase    ,   )
                      }
                 records.append(__lowercase   )
            return records






@add_end_docstrings(A )
class 	lowerCAmelCase  (   A ):
       lowerCAmelCase_												=			"summary"







       def __call__(				self      :			Dict    ,   *__lowercase      :			Dict    ,   **__lowercase      :			Dict   ):







            """simple docstring"""







            return super().__call__(*__lowercase    ,   **__lowercase   )







       def       snake_case (				self      :			str    ,   __lowercase      :			int    ,   __lowercase      :			int    ,   __lowercase      :			int   ):







            """simple docstring"""







            if max_length < min_length:
                 logger.warning(f'''Your min_length={min_length} must be inferior than your max_length={max_length}.'''   )

            if input_length < max_length:
                 logger.warning(
                     f'''Your max_length is set to {max_length}, but your input_length is only {input_length}. Since this is '''
                     'a summarization task, where outputs shorter than the input are typically wanted, you might '
                     f'''consider decreasing max_length manually, e.g. summarizer(\'...\', max_length={input_length//2})'''   )






@add_end_docstrings(A )
class 	lowerCAmelCase  (   A ):
       lowerCAmelCase_												=			"translation"







       def       snake_case (				self      :			int    ,   __lowercase      :			int    ,   __lowercase      :			int    ,   __lowercase      :			int   ):







            """simple docstring"""







            if input_length > 0.9 * max_length:
                 logger.warning(
                     f'''Your input_length: {input_length} is bigger than 0.9 * max_length: {max_length}. You might consider '''
                     'increasing your max_length manually, e.g. translator(\'...\', max_length=400)'   )
            return True







       def       snake_case (				self      :			int    ,   *__lowercase      :			str    ,   __lowercase      :			List[str]=TruncationStrategy.DO_NOT_TRUNCATE    ,   __lowercase      :			List[Any]=None    ,   __lowercase      :			Optional[int]=None   ):







            """simple docstring"""







            if getattr(self.tokenizer    ,   '_build_translation_inputs'    ,   __lowercase   ):
                 return self.tokenizer._build_translation_inputs(
                     *__lowercase    ,   return_tensors=self.framework    ,   truncation=__lowercase    ,   src_lang=__lowercase    ,   tgt_lang=__lowercase   )
            else:
                 return super()._parse_and_tokenize(*__lowercase    ,   truncation=__lowercase   )







       def       snake_case (				self      :			Optional[int]    ,   __lowercase      :			List[str]=None    ,   __lowercase      :			Optional[Any]=None    ,   **__lowercase      :			int   ):







            """simple docstring"""







            __lowercase		,  __lowercase		,  __lowercase      		=super()._sanitize_parameters(**__lowercase   )
            if src_lang is not None:
                 __lowercase      		=src_lang
            if tgt_lang is not None:
                 __lowercase      		=tgt_lang
            if src_lang is None and tgt_lang is None:
                 # Backward compatibility, direct arguments use is preferred.
                 __lowercase      		=kwargs.get('task'    ,   self.task   )
                 __lowercase      		=task.split('_'   )
                 if task and len(__lowercase   ) == 4:
                      # translation, XX, to YY
                      __lowercase      		=items[1]
                      __lowercase      		=items[3]
            return preprocess_params, forward_params, postprocess_params







       def __call__(				self      :			Optional[Any]    ,   *__lowercase      :			Any    ,   **__lowercase      :			List[Any]   ):







            """simple docstring"""







            return super().__call__(*__lowercase    ,   **__lowercase   )


