





from typing import Optional, Tuple

import flax.linen as nn
import jax
import jax.numpy as jnp
from configuration_hybrid_clip import HybridCLIPConfig
from flax.core.frozen_dict import FrozenDict

from transformers import FLAX_MODEL_MAPPING, FlaxCLIPVisionModel
from transformers.modeling_flax_utils import FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPOutput
from transformers.utils import logging


_UpperCAmelCase			   =							logging.get_logger(__name__)






class    snake_case_				(		nn.Module						):
  A_										=    42
  A_										=    jnp.floataa






  def       UpperCAmelCase__						(							self :     Dict			)->str:






   '''simple docstring'''

   __lowerCAmelCase     :       Union[str, Any]						=   self.config.text_config
   __lowerCAmelCase     :       Union[str, Any]						=   self.config.vision_config

   __lowerCAmelCase     :       Optional[Any]						=   self.config.projection_dim
   __lowerCAmelCase     :       int						=   text_config.hidden_size
   __lowerCAmelCase     :       Tuple						=   vision_config.hidden_size

   __lowerCAmelCase     :       Dict						=   FLAX_MODEL_MAPPING[self.config.text_config.__class__].module_class
   __lowerCAmelCase     :       str						=   FLAX_MODEL_MAPPING.get(self.config.vision_config.__class__   ,  _UpperCamelCase			).module_class

   __lowerCAmelCase     :       List[str]						=   text_module(_UpperCamelCase   ,  dtype=self.dtype			)
   __lowerCAmelCase     :       Dict						=   vision_module(_UpperCamelCase   ,  dtype=self.dtype			)

   __lowerCAmelCase     :       List[str]						=   nn.Dense(
       self.projection_dim   ,  dtype=self.dtype   ,  kernel_init=jax.nn.initializers.normal(0.02			)   ,  use_bias=_UpperCamelCase   ,  )
   __lowerCAmelCase     :       int						=   nn.Dense(
       self.projection_dim   ,  dtype=self.dtype   ,  kernel_init=jax.nn.initializers.normal(0.02			)   ,  use_bias=_UpperCamelCase   ,  )
   __lowerCAmelCase     :       Dict						=   self.param("""logit_scale"""   ,  jax.nn.initializers.ones   ,  []			)






  def __call__(							self :     Dict   ,  _snake_case :     Dict=None   ,  _snake_case :     Tuple=None   ,  _snake_case :     List[Any]=None   ,  _snake_case :     Optional[Any]=None   ,  _snake_case :     Optional[Any]=None   ,  _snake_case :     Optional[int] = True   ,  _snake_case :     str=None   ,  _snake_case :     Dict=None   ,  _snake_case :     List[Any]=None   ,  )->int:






   '''simple docstring'''

   __lowerCAmelCase     :       str						=   return_dict if return_dict is not None else self.config.return_dict

   __lowerCAmelCase     :       str						=   self.vision_model(
       pixel_values=_UpperCamelCase   ,  deterministic=_UpperCamelCase   ,  output_attentions=_UpperCamelCase   ,  output_hidden_states=_UpperCamelCase   ,  return_dict=_UpperCamelCase   ,  )

   __lowerCAmelCase     :       Optional[int]						=   self.text_model(
       input_ids=_UpperCamelCase   ,  attention_mask=_UpperCamelCase   ,  token_type_ids=_UpperCamelCase   ,  position_ids=_UpperCamelCase   ,  deterministic=_UpperCamelCase   ,  output_attentions=_UpperCamelCase   ,  output_hidden_states=_UpperCamelCase   ,  return_dict=_UpperCamelCase   ,  )

   __lowerCAmelCase     :       str						=   vision_outputs[1]
   __lowerCAmelCase     :       List[Any]						=   self.visual_projection(_UpperCamelCase			)

   __lowerCAmelCase     :       int						=   text_outputs[1]
   __lowerCAmelCase     :       Union[str, Any]						=   self.text_projection(_UpperCamelCase			)

   # normalized features
   __lowerCAmelCase     :       Union[str, Any]						=   image_embeds / jnp.linalg.norm(_UpperCamelCase   ,  axis=-1   ,  keepdims=_UpperCamelCase			)
   __lowerCAmelCase     :       Any						=   text_embeds / jnp.linalg.norm(_UpperCamelCase   ,  axis=-1   ,  keepdims=_UpperCamelCase			)

   # cosine similarity as logits
   __lowerCAmelCase     :       str						=   jnp.exp(self.logit_scale			)
   __lowerCAmelCase     :       Optional[int]						=   jnp.matmul(_UpperCamelCase   ,  image_embeds.T			) * logit_scale
   __lowerCAmelCase     :       List[Any]						=   logits_per_text.T

   if not return_dict:
    return (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)

   return FlaxCLIPOutput(
       logits_per_image=_UpperCamelCase   ,  logits_per_text=_UpperCamelCase   ,  text_embeds=_UpperCamelCase   ,  image_embeds=_UpperCamelCase   ,  text_model_output=_UpperCamelCase   ,  vision_model_output=_UpperCamelCase   ,  )






class    snake_case_				(		__lowercase						):
  A_										=    HybridCLIPConfig
  A_										=    FlaxHybridCLIPModule






  def __init__(							self :     Dict   ,  _snake_case :     List[Any]   ,  _snake_case :     Optional[Any] = None   ,  _snake_case :     Tuple = 0   ,  _snake_case :     List[Any] = jnp.floataa   ,  **_snake_case :     List[Any]   ,  )->str:






   '''simple docstring'''

   if input_shape is None:
    __lowerCAmelCase     :       Optional[int]						=   ((1, 1), (1, config.vision_config.image_size, config.vision_config.image_size, 3))

   __lowerCAmelCase     :       Dict						=   self.module_class(config=_UpperCamelCase   ,  dtype=_UpperCamelCase   ,  **_UpperCamelCase			)
   super().__init__(_UpperCamelCase   ,  _UpperCamelCase   ,  input_shape=_UpperCamelCase   ,  seed=_UpperCamelCase   ,  dtype=_UpperCamelCase			)






  def       UpperCAmelCase__						(							self :     Optional[int]   ,  _snake_case :     Any   ,  _snake_case :     List[Any]   ,  _snake_case :     str = None			)->Dict:






   '''simple docstring'''

   __lowerCAmelCase     :       List[str]						=   jnp.zeros(input_shape[0]   ,  dtype="""i4"""			)
   __lowerCAmelCase     :       Optional[Any]						=   jnp.broadcast_to(jnp.arange(jnp.atleast_ad(_UpperCamelCase			).shape[-1]			)   ,  input_shape[0]			)
   __lowerCAmelCase     :       str						=   jnp.ones_like(_UpperCamelCase			)
   __lowerCAmelCase     :       Any						=   jnp.ones_like(_UpperCamelCase			)

   __lowerCAmelCase     :       Union[str, Any]						=   jax.random.normal(_UpperCamelCase   ,  input_shape[1]			)

   __lowerCAmelCase					,     __lowerCAmelCase     :       Tuple						=   jax.random.split(_UpperCamelCase			)
   __lowerCAmelCase     :       str						=   {"""params""": params_rng, """dropout""": dropout_rng}

   return self.module.init(_UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase			)["params"]






  def __call__(							self :     Union[str, Any]   ,  _snake_case :     Any   ,  _snake_case :     int   ,  _snake_case :     List[Any]=None   ,  _snake_case :     str=None   ,  _snake_case :     List[str]=None   ,  _snake_case :     str = None   ,  _snake_case :     str = None   ,  _snake_case :     Union[str, Any] = False   ,  _snake_case :     List[str] = None   ,  _snake_case :     int = None   ,  _snake_case :     Any = None   ,  )->Any:






   '''simple docstring'''

   __lowerCAmelCase     :       int						=   output_attentions if output_attentions is not None else self.config.output_attentions
   __lowerCAmelCase     :       List[str]						=   (
       output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
   )
   __lowerCAmelCase     :       Union[str, Any]						=   return_dict if return_dict is not None else self.config.return_dict

   if position_ids is None:
    __lowerCAmelCase     :       Union[str, Any]						=   jnp.broadcast_to(jnp.arange(jnp.atleast_ad(_UpperCamelCase			).shape[-1]			)   ,  input_ids.shape			)

   if token_type_ids is None:
    __lowerCAmelCase     :       Any						=   jnp.zeros_like(_UpperCamelCase			)

   if attention_mask is None:
    __lowerCAmelCase     :       Union[str, Any]						=   jnp.ones_like(_UpperCamelCase			)

   # Handle any PRNG if needed
   __lowerCAmelCase     :       Optional[Any]						=   {}
   if dropout_rng is not None:
    __lowerCAmelCase     :       Tuple						=   dropout_rng

   return self.module.apply(
       {"""params""": params or self.params}   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  jnp.array(_UpperCamelCase   ,  dtype=jnp.floataa			)   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  not train   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  rngs=_UpperCamelCase   ,  )






  def       UpperCAmelCase__						(							self :     Union[str, Any]   ,  _snake_case :     Tuple   ,  _snake_case :     int=None   ,  _snake_case :     Any=None   ,  _snake_case :     str=None   ,  _snake_case :     Optional[Any] = None   ,  _snake_case :     Optional[int] = None   ,  _snake_case :     Optional[Any]=False   ,  )->Optional[Any]:






   '''simple docstring'''

   if position_ids is None:
    __lowerCAmelCase     :       Optional[Any]						=   jnp.broadcast_to(jnp.arange(jnp.atleast_ad(_UpperCamelCase			).shape[-1]			)   ,  input_ids.shape			)

   if token_type_ids is None:
    __lowerCAmelCase     :       List[Any]						=   jnp.zeros_like(_UpperCamelCase			)

   if attention_mask is None:
    __lowerCAmelCase     :       Optional[Any]						=   jnp.ones_like(_UpperCamelCase			)

   # Handle any PRNG if needed
   __lowerCAmelCase     :       Optional[int]						=   {}
   if dropout_rng is not None:
    __lowerCAmelCase     :       Any						=   dropout_rng

   def _get_features(_snake_case :     int   ,  _snake_case :     Dict   ,  _snake_case :     Tuple   ,  _snake_case :     Optional[int]   ,  _snake_case :     int   ,  _snake_case :     Optional[int]			):
    __lowerCAmelCase     :       List[Any]						=   module.text_model(
        input_ids=_UpperCamelCase   ,  attention_mask=_UpperCamelCase   ,  position_ids=_UpperCamelCase   ,  token_type_ids=_UpperCamelCase   ,  deterministic=_UpperCamelCase   ,  )
    __lowerCAmelCase     :       List[Any]						=   text_outputs[1]
    __lowerCAmelCase     :       Dict						=   module.text_projection(_UpperCamelCase			)
    return text_features

   return self.module.apply(
       {"""params""": params or self.params}   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  jnp.array(_UpperCamelCase   ,  dtype="""i4"""			)   ,  not train   ,  method=_get_features   ,  rngs=_UpperCamelCase   ,  )






  def       UpperCAmelCase__						(							self :     str   ,  _snake_case :     Union[str, Any]   ,  _snake_case :     List[str] = None   ,  _snake_case :     Optional[Any] = None   ,  _snake_case :     List[str]=False			)->List[Any]:






   '''simple docstring'''

   __lowerCAmelCase     :       List[Any]						=   {}
   if dropout_rng is not None:
    __lowerCAmelCase     :       Union[str, Any]						=   dropout_rng

   def _get_features(_snake_case :     str   ,  _snake_case :     str   ,  _snake_case :     int			):
    __lowerCAmelCase     :       str						=   module.vision_model(pixel_values=_UpperCamelCase   ,  deterministic=_UpperCamelCase			)
    __lowerCAmelCase     :       int						=   vision_outputs[1]  # pooled_output
    __lowerCAmelCase     :       str						=   module.visual_projection(_UpperCamelCase			)
    return image_features

   return self.module.apply(
       {"""params""": params or self.params}   ,  jnp.array(_UpperCamelCase   ,  dtype=jnp.floataa			)   ,  not train   ,  method=_get_features   ,  rngs=_UpperCamelCase   ,  )






  @classmethod
  def       UpperCAmelCase__						(							cls :     List[str]   ,  _snake_case :     Optional[int] = None   ,  _snake_case :     str = None   ,  *_snake_case :     int   ,  **_snake_case :     int   ,  )->List[Any]:






   '''simple docstring'''

   __lowerCAmelCase     :       Dict						=   {
       argument[len("""text_"""			) :]: value for argument, value in kwargs.items() if argument.startswith("""text_"""			)
   }

   __lowerCAmelCase     :       int						=   {
       argument[len("""vision_"""			) :]: value for argument, value in kwargs.items() if argument.startswith("""vision_"""			)
   }

   # remove text, vision kwargs from kwargs
   for key in kwargs_text.keys():
    del kwargs["text_" + key]
   for key in kwargs_vision.keys():
    del kwargs["vision_" + key]

   # Load and initialize the text and vision model
   __lowerCAmelCase     :       List[str]						=   kwargs_text.pop("""model"""   ,  _UpperCamelCase			)
   if text_model is None:
    assert (
        text_model_name_or_path is not None
    ), "If `model` is not defined as an argument, a `text_model_name_or_path` has to be defined"
    from transformers import FlaxAutoModel

    if "config" not in kwargs_text:
     from transformers import AutoConfig

     __lowerCAmelCase     :       Optional[int]						=   AutoConfig.from_pretrained(_UpperCamelCase			)
     __lowerCAmelCase     :       Any						=   text_config

    __lowerCAmelCase     :       Union[str, Any]						=   FlaxAutoModel.from_pretrained(_UpperCamelCase   ,  *_UpperCamelCase   ,  **_UpperCamelCase			)

   __lowerCAmelCase     :       Tuple						=   kwargs_vision.pop("""model"""   ,  _UpperCamelCase			)
   if vision_model is None:
    assert (
        vision_model_name_or_path is not None
    ), "If `model` is not defined as an argument, a `vision_model_name_or_path` has to be defined"
    from transformers import FlaxAutoModel

    if "config" not in kwargs_vision:
     from transformers import AutoConfig

     __lowerCAmelCase     :       List[str]						=   AutoConfig.from_pretrained(_UpperCamelCase			)
     __lowerCAmelCase     :       Any						=   vision_config

    __lowerCAmelCase     :       Optional[Any]						=   FlaxAutoModel.from_pretrained(_UpperCamelCase   ,  *_UpperCamelCase   ,  **_UpperCamelCase			)

   # instantiate config with corresponding kwargs
   __lowerCAmelCase     :       Any						=   kwargs.pop("""dtype"""   ,  jnp.floataa			)
   __lowerCAmelCase     :       Union[str, Any]						=   HybridCLIPConfig.from_text_vision_configs(text_model.config   ,  vision_model.config   ,  **_UpperCamelCase			)

   # init model
   __lowerCAmelCase     :       Dict						=   cls(_UpperCamelCase   ,  *_UpperCamelCase   ,  dtype=_UpperCamelCase   ,  **_UpperCamelCase			)

   if vision_config.model_type == "clip":
    __lowerCAmelCase     :       Any						=   vision_model.params["""vision_model"""]
    __lowerCAmelCase     :       Tuple						=   vision_model.params["""visual_projection"""]["""kernel"""]
   else:
    __lowerCAmelCase     :       Optional[int]						=   vision_model.params

   __lowerCAmelCase     :       List[Any]						=   text_model.params

   return model







import csv
import dataclasses
import json
from dataclasses import dataclass
from typing import List, Optional, Union

from ...utils import is_tf_available, is_torch_available, logging


__snake_case			:			Union[str, Any]			   =   logging.get_logger(__name__)


@dataclass
class 			__SCREAMING_SNAKE_CASE    :
						_SCREAMING_SNAKE_CASE		:				str
						_SCREAMING_SNAKE_CASE		:				str
						_SCREAMING_SNAKE_CASE		:				Optional[str]            =    None
						_SCREAMING_SNAKE_CASE		:				Optional[str]            =    None






						def 					UpperCamelCase__					(		self      ):






												"""simple docstring"""
												return json.dumps(dataclasses.asdict(self      )   ,  indent=2      ) + "\n"




@dataclass(frozen=__lowercase)
class 			__SCREAMING_SNAKE_CASE    :
						_SCREAMING_SNAKE_CASE		:				List[int]
						_SCREAMING_SNAKE_CASE		:				Optional[List[int]]            =    None
						_SCREAMING_SNAKE_CASE		:				Optional[List[int]]            =    None
						_SCREAMING_SNAKE_CASE		:				Optional[Union[int, float]]            =    None






						def 					UpperCamelCase__					(		self      ):






												"""simple docstring"""
												return json.dumps(dataclasses.asdict(self      )      ) + "\n"




class 			__SCREAMING_SNAKE_CASE    :






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase      ):






												"""simple docstring"""
												raise NotImplementedError()






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase      ):






												"""simple docstring"""
												raise NotImplementedError()






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase      ):






												"""simple docstring"""
												raise NotImplementedError()






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase      ):






												"""simple docstring"""
												raise NotImplementedError()






						def 					UpperCamelCase__					(		self      ):






												"""simple docstring"""
												raise NotImplementedError()






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase      ):






												"""simple docstring"""
												if len(self.get_labels()      ) > 1:
																		lowerCAmelCase__			  =  self.get_labels()[int(example.label      )]
												return example







						@classmethod
						def 					UpperCamelCase__					(		cls   ,  _UpperCamelCase   ,  _UpperCamelCase=None      ):






												"""simple docstring"""
												with open(_UpperCamelCase   ,  'r'   ,  encoding='utf-8-sig'      ) as f:
																		return list(csv.reader(_UpperCamelCase   ,  delimiter='\t'   ,  quotechar=_UpperCamelCase      )      )




class 			__SCREAMING_SNAKE_CASE    (      __lowercase):






						def __init__(		self   ,  _UpperCamelCase=None   ,  _UpperCamelCase=None   ,  _UpperCamelCase="classification"   ,  _UpperCamelCase=False      ):






												"""simple docstring"""
												lowerCAmelCase__			  =  [] if labels is None else labels
												lowerCAmelCase__			  =  [] if examples is None else examples
												lowerCAmelCase__			  =  mode
												lowerCAmelCase__			  =  verbose






						def __len__(		self      ):






												"""simple docstring"""
												return len(self.examples      )






						def __getitem__(		self   ,  _UpperCamelCase      ):






												"""simple docstring"""
												if isinstance(_UpperCamelCase   ,  _UpperCamelCase      ):
																		return SingleSentenceClassificationProcessor(labels=self.labels   ,  examples=self.examples[idx]      )
												return self.examples[idx]






						@classmethod
						def 					UpperCamelCase__					(		cls   ,  _UpperCamelCase   ,  _UpperCamelCase=""   ,  _UpperCamelCase=0   ,  _UpperCamelCase=1   ,  _UpperCamelCase=None   ,  _UpperCamelCase=False   ,  **_UpperCamelCase      ):






												"""simple docstring"""
												lowerCAmelCase__			  =  cls(**_UpperCamelCase      )
												processor.add_examples_from_csv(
												    _UpperCamelCase   ,  split_name=_UpperCamelCase   ,  column_label=_UpperCamelCase   ,  column_text=_UpperCamelCase   ,  column_id=_UpperCamelCase   ,  skip_first_row=_UpperCamelCase   ,  overwrite_labels=_UpperCamelCase   ,  overwrite_examples=_UpperCamelCase   ,  )
												return processor






						@classmethod
						def 					UpperCamelCase__					(		cls   ,  _UpperCamelCase   ,  _UpperCamelCase=None   ,  **_UpperCamelCase      ):






												"""simple docstring"""
												lowerCAmelCase__			  =  cls(**_UpperCamelCase      )
												processor.add_examples(_UpperCamelCase   ,  labels=_UpperCamelCase      )
												return processor






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase=""   ,  _UpperCamelCase=0   ,  _UpperCamelCase=1   ,  _UpperCamelCase=None   ,  _UpperCamelCase=False   ,  _UpperCamelCase=False   ,  _UpperCamelCase=False   ,  ):






												"""simple docstring"""
												lowerCAmelCase__			  =  self._read_tsv(_UpperCamelCase      )
												if skip_first_row:
																		lowerCAmelCase__			  =  lines[1:]
												lowerCAmelCase__			  =  []
												lowerCAmelCase__			  =  []
												lowerCAmelCase__			  =  []
												for i, line in enumerate(_UpperCamelCase      ):
																		texts.append(line[column_text]      )
																		labels.append(line[column_label]      )
																		if column_id is not None:
																								ids.append(line[column_id]      )
																		else:
																								lowerCAmelCase__			  =  F"{split_name}-{i}" if split_name else str(_UpperCamelCase      )
																								ids.append(_UpperCamelCase      )

												return self.add_examples(
												    _UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase   ,  overwrite_labels=_UpperCamelCase   ,  overwrite_examples=_UpperCamelCase      )






						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase=None   ,  _UpperCamelCase=None   ,  _UpperCamelCase=False   ,  _UpperCamelCase=False      ):






												"""simple docstring"""
												if labels is not None and len(_UpperCamelCase      ) != len(_UpperCamelCase      ):
																		raise ValueError(
																		    F"Text and labels have mismatched lengths {len(_UpperCamelCase      )} and {len(_UpperCamelCase      )}"      )
												if ids is not None and len(_UpperCamelCase      ) != len(_UpperCamelCase      ):
																		raise ValueError(F"Text and ids have mismatched lengths {len(_UpperCamelCase      )} and {len(_UpperCamelCase      )}"      )
												if ids is None:
																		lowerCAmelCase__			  =  [None] * len(_UpperCamelCase      )
												if labels is None:
																		lowerCAmelCase__			  =  [None] * len(_UpperCamelCase      )
												lowerCAmelCase__			  =  []
												lowerCAmelCase__			  =  set()
												for text_or_text_and_label, label, guid in zip(_UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase      ):
																		if isinstance(_UpperCamelCase   ,  (tuple, list)      ) and label is None:
																								lowerCAmelCase__  ,				lowerCAmelCase__			  =  text_or_text_and_label
																		else:
																								lowerCAmelCase__			  =  text_or_text_and_label
																		added_labels.add(_UpperCamelCase      )
																		examples.append(InputExample(guid=_UpperCamelCase   ,  text_a=_UpperCamelCase   ,  text_b=_UpperCamelCase   ,  label=_UpperCamelCase      )      )

												# Update examples
												if overwrite_examples:
																		lowerCAmelCase__			  =  examples
												else:
																		self.examples.extend(_UpperCamelCase      )

												# Update labels
												if overwrite_labels:
																		lowerCAmelCase__			  =  list(_UpperCamelCase      )
												else:
																		lowerCAmelCase__			  =  list(set(self.labels      ).union(_UpperCamelCase      )      )

												return self.examples







						def 					UpperCamelCase__					(		self   ,  _UpperCamelCase   ,  _UpperCamelCase=None   ,  _UpperCamelCase=False   ,  _UpperCamelCase=0   ,  _UpperCamelCase=True   ,  _UpperCamelCase=None   ,  ):






												"""simple docstring"""
												if max_length is None:
																		lowerCAmelCase__			  =  tokenizer.max_len

												lowerCAmelCase__			  =  {label: i for i, label in enumerate(self.labels      )}

												lowerCAmelCase__			  =  []
												for ex_index, example in enumerate(self.examples      ):
																		if ex_index % 1_00_00 == 0:
																								logger.info(F"Tokenizing example {ex_index}"      )

																		lowerCAmelCase__			  =  tokenizer.encode(
																		    example.text_a   ,  add_special_tokens=_UpperCamelCase   ,  max_length=min(_UpperCamelCase   ,  tokenizer.max_len      )   ,  )
																		all_input_ids.append(_UpperCamelCase      )

												lowerCAmelCase__			  =  max(len(_UpperCamelCase      ) for input_ids in all_input_ids      )

												lowerCAmelCase__			  =  []
												for ex_index, (input_ids, example) in enumerate(zip(_UpperCamelCase   ,  self.examples      )      ):
																		if ex_index % 1_00_00 == 0:
																								logger.info(F"Writing example {ex_index}/{len(self.examples      )}"      )
																		# The mask has 1 for real tokens and 0 for padding tokens. Only real
																		# tokens are attended to.
																		lowerCAmelCase__			  =  [1 if mask_padding_with_zero else 0] * len(_UpperCamelCase      )

																		# Zero-pad up to the sequence length.
																		lowerCAmelCase__			  =  batch_length - len(_UpperCamelCase      )
																		if pad_on_left:
																								lowerCAmelCase__			  =  ([pad_token] * padding_length) + input_ids
																								lowerCAmelCase__			  =  ([0 if mask_padding_with_zero else 1] * padding_length) + attention_mask
																		else:
																								lowerCAmelCase__			  =  input_ids + ([pad_token] * padding_length)
																								lowerCAmelCase__			  =  attention_mask + ([0 if mask_padding_with_zero else 1] * padding_length)

																		if len(_UpperCamelCase      ) != batch_length:
																								raise ValueError(F"Error with input length {len(_UpperCamelCase      )} vs {batch_length}"      )
																		if len(_UpperCamelCase      ) != batch_length:
																								raise ValueError(F"Error with input length {len(_UpperCamelCase      )} vs {batch_length}"      )

																		if self.mode == "classification":
																								lowerCAmelCase__			  =  label_map[example.label]
																		elif self.mode == "regression":
																								lowerCAmelCase__			  =  float(example.label      )
																		else:
																								raise ValueError(self.mode      )

																		if ex_index < 5 and self.verbose:
																								logger.info('*** Example ***'      )
																								logger.info(F"guid: {example.guid}"      )
																								logger.info(F"input_ids: {' '.join([str(_UpperCamelCase      ) for x in input_ids]      )}"      )
																								logger.info(F"attention_mask: {' '.join([str(_UpperCamelCase      ) for x in attention_mask]      )}"      )
																								logger.info(F"label: {example.label} (id = {label})"      )

																		features.append(InputFeatures(input_ids=_UpperCamelCase   ,  attention_mask=_UpperCamelCase   ,  label=_UpperCamelCase      )      )

												if return_tensors is None:
																		return features
												elif return_tensors == "tf":
																		if not is_tf_available():
																								raise RuntimeError('return_tensors set to \'tf\' but TensorFlow 2.0 can\'t be imported'      )
																		import tensorflow as tf

																		def gen():
																								for ex in features:
																														yield ({"input_ids": ex.input_ids, "attention_mask": ex.attention_mask}, ex.label)

																		lowerCAmelCase__			  =  tf.data.Dataset.from_generator(
																		    _UpperCamelCase   ,  ({'input_ids': tf.intaa, 'attention_mask': tf.intaa}, tf.intaa)   ,  ({'input_ids': tf.TensorShape([None]      ), 'attention_mask': tf.TensorShape([None]      )}, tf.TensorShape([]      ))   ,  )
																		return dataset
												elif return_tensors == "pt":
																		if not is_torch_available():
																								raise RuntimeError('return_tensors set to \'pt\' but PyTorch can\'t be imported'      )
																		import torch
																		from torch.utils.data import TensorDataset

																		lowerCAmelCase__			  =  torch.tensor([f.input_ids for f in features]   ,  dtype=torch.long      )
																		lowerCAmelCase__			  =  torch.tensor([f.attention_mask for f in features]   ,  dtype=torch.long      )
																		if self.mode == "classification":
																								lowerCAmelCase__			  =  torch.tensor([f.label for f in features]   ,  dtype=torch.long      )
																		elif self.mode == "regression":
																								lowerCAmelCase__			  =  torch.tensor([f.label for f in features]   ,  dtype=torch.float      )

																		lowerCAmelCase__			  =  TensorDataset(_UpperCamelCase   ,  _UpperCamelCase   ,  _UpperCamelCase      )
																		return dataset
												else:
																		raise ValueError('return_tensors should be one of \'tf\' or \'pt\''      )





