


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 ):

 '''simple docstring'''

 __snake_case				=						0
 __snake_case				=						1







@add_end_docstrings(lowerCamelCase__ )
class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						"generated"



 def __init__(			self	,   *_UpperCAmelCase	,   **_UpperCAmelCase      ):
       super().__init__(*_UpperCAmelCase	,   **_UpperCAmelCase      )

       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 		UpperCamelCase__  (			self	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None	,   **_UpperCAmelCase	,   ):
       snake_case_ =   {}
       if truncation is not None:
             snake_case_ =   truncation

       snake_case_ =   generate_kwargs

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

       if clean_up_tokenization_spaces is not None:
             snake_case_ =   clean_up_tokenization_spaces

       if stop_sequence is not None:
             snake_case_ =   self.tokenizer.encode(_UpperCAmelCase	,   add_special_tokens=_UpperCAmelCase      )
             if len(_UpperCAmelCase      ) > 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.'''      )
             snake_case_ =   stop_sequence_ids[0]

       return preprocess_params, forward_params, postprocess_params



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       return True



 def 		UpperCamelCase__  (			self	,   *_UpperCAmelCase	,   _UpperCAmelCase      ):
       snake_case_ =   self.model.config.prefix if self.model.config.prefix is not None else ''''''
       if isinstance(args[0]	,   _UpperCAmelCase      ):
             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'''      )
             snake_case_ =   ([prefix + arg for arg in args[0]],)
             snake_case_ =   True

       elif isinstance(args[0]	,   _UpperCAmelCase      ):
             snake_case_ =   (prefix + args[0],)
             snake_case_ =   False
       else:
             raise ValueError(
                 F''' `args[0]`: {args[0]} have the wrong format. The should be either of type `str` or type `list`'''      )
       snake_case_ =   self.tokenizer(*_UpperCAmelCase	,   padding=_UpperCAmelCase	,   truncation=_UpperCAmelCase	,   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	,   *_UpperCAmelCase	,   **_UpperCAmelCase      ):
       snake_case_ =   super().__call__(*_UpperCAmelCase	,   **_UpperCAmelCase      )
       if (
           isinstance(args[0]	,   _UpperCAmelCase      )
           and all(isinstance(_UpperCAmelCase	,   _UpperCAmelCase      ) for el in args[0]      )
           and all(len(_UpperCAmelCase      ) == 1 for res in result      )
       ):
             return [res[0] for res in result]
       return result



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase=TruncationStrategy.DO_NOT_TRUNCATE	,   **_UpperCAmelCase      ):
       snake_case_ =   self._parse_and_tokenize(_UpperCAmelCase	,   truncation=_UpperCAmelCase	,   **_UpperCAmelCase      )
       return inputs



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   **_UpperCAmelCase      ):
       if self.framework == "pt":
             snake_case_ ,					snake_case_ =   model_inputs['''input_ids'''].shape
       elif self.framework == "tf":
             snake_case_ ,					snake_case_ =   tf.shape(model_inputs['''input_ids''']      ).numpy()

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





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







@add_end_docstrings(lowerCamelCase__ )
class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						"summary"



 def __call__(			self	,   *_UpperCAmelCase	,   **_UpperCAmelCase      ):
       return super().__call__(*_UpperCAmelCase	,   **_UpperCAmelCase      )





 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       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(lowerCamelCase__ )
class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						"translation"



 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase	,   _UpperCAmelCase	,   _UpperCAmelCase      ):
       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 		UpperCamelCase__  (			self	,   *_UpperCAmelCase	,   _UpperCAmelCase=TruncationStrategy.DO_NOT_TRUNCATE	,   _UpperCAmelCase=None	,   _UpperCAmelCase=None      ):
       if getattr(self.tokenizer	,   '''_build_translation_inputs'''	,   _UpperCAmelCase      ):
             return self.tokenizer._build_translation_inputs(
                 *_UpperCAmelCase	,   return_tensors=self.framework	,   truncation=_UpperCAmelCase	,   src_lang=_UpperCAmelCase	,   tgt_lang=_UpperCAmelCase      )
       else:
             return super()._parse_and_tokenize(*_UpperCAmelCase	,   truncation=_UpperCAmelCase      )



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





 def __call__(			self	,   *_UpperCAmelCase	,   **_UpperCAmelCase      ):
       return super().__call__(*_UpperCAmelCase	,   **_UpperCAmelCase      )



from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput







@dataclass
class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						None
 __snake_case				=						None
 __snake_case				=						None
 __snake_case				=						None







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 def __init__(			self	,   _UpperCAmelCase=1	,   _UpperCAmelCase=0	,   _UpperCAmelCase=2	,   _UpperCAmelCase=5_12	,   _UpperCAmelCase="cls"	,   _UpperCAmelCase=False	,   _UpperCAmelCase=True	,   **_UpperCAmelCase	,   ):
       super().__init__(pad_token_id=_UpperCAmelCase	,   bos_token_id=_UpperCAmelCase	,   eos_token_id=_UpperCAmelCase	,   **_UpperCAmelCase      )
       snake_case_ =   project_dim
       snake_case_ =   pooler_fn
       snake_case_ =   learn_encoder
       snake_case_ =   use_attention_mask







class lowerCAmelCase_		(				lowerCamelCase__ ):

 '''simple docstring'''

 __snake_case				=						[R"pooler", R"logit_scale"]
 __snake_case				=						[R"position_ids", R"predictions.decoder.bias"]
 __snake_case				=						"roberta"
 __snake_case				=						RobertaSeriesConfig



 def __init__(			self	,   _UpperCAmelCase      ):
       super().__init__(_UpperCAmelCase      )
       snake_case_ =   XLMRobertaModel(_UpperCAmelCase      )
       snake_case_ =   nn.Linear(config.hidden_size	,   config.project_dim      )
       snake_case_ =   getattr(_UpperCAmelCase	,   '''has_pre_transformation'''	,   _UpperCAmelCase      )
       if self.has_pre_transformation:
             snake_case_ =   nn.Linear(config.hidden_size	,   config.project_dim      )
             snake_case_ =   nn.LayerNorm(config.hidden_size	,   eps=config.layer_norm_eps      )
       self.post_init()





 def 		UpperCamelCase__  (			self	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   _UpperCAmelCase = None	,   ):
       snake_case_ =   return_dict if return_dict is not None else self.config.use_return_dict

       snake_case_ =   self.base_model(
           input_ids=_UpperCAmelCase	,   attention_mask=_UpperCAmelCase	,   token_type_ids=_UpperCAmelCase	,   position_ids=_UpperCAmelCase	,   head_mask=_UpperCAmelCase	,   inputs_embeds=_UpperCAmelCase	,   encoder_hidden_states=_UpperCAmelCase	,   encoder_attention_mask=_UpperCAmelCase	,   output_attentions=_UpperCAmelCase	,   output_hidden_states=True if self.has_pre_transformation else output_hidden_states	,   return_dict=_UpperCAmelCase	,   )

       if self.has_pre_transformation:
             snake_case_ =   outputs['''hidden_states'''][-2]
             snake_case_ =   self.pre_LN(_UpperCAmelCase      )
             snake_case_ =   self.transformation_pre(_UpperCAmelCase      )

             return TransformationModelOutput(
                 projection_state=_UpperCAmelCase	,   last_hidden_state=outputs.last_hidden_state	,   hidden_states=outputs.hidden_states	,   attentions=outputs.attentions	,   )
       else:
             snake_case_ =   self.transformation(outputs.last_hidden_state      )
             return TransformationModelOutput(
                 projection_state=_UpperCAmelCase	,   last_hidden_state=outputs.last_hidden_state	,   hidden_states=outputs.hidden_states	,   attentions=outputs.attentions	,   )