


'''simple docstring'''






# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch

from ..models.auto import AutoModelForSequenceClassification, AutoTokenizer
from .base import PipelineTool




class 						a__							(				a__    ):





							'''simple docstring'''






							lowercase__     :  List[str]      		= "facebook/bart-large-mnli"
							lowercase__     :  str      		= (
							    "This is a tool that classifies an English text using provided labels. It takes two inputs: `text`, which "
							    "should be the text to classify, and `labels`, which should be the list of labels to use for classification. "
							    "It returns the most likely label in the list of provided `labels` for the input text."
							)
							lowercase__     :  List[str]      		= "text_classifier"
							lowercase__     :  List[str]      		= AutoTokenizer
							lowercase__     :  int      		= AutoModelForSequenceClassification

							lowercase__     :  Optional[Any]      		= ["text", ["text"]]
							lowercase__     :  str      		= ["text"]







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[Any]:
														super().setup()
														lowerCAmelCase__			  =      self.model.config
														lowerCAmelCase__			  =      -1
														for idx, label in config.idalabel.items():
																					if label.lower().startswith('''entail'''  ):
																												lowerCAmelCase__			  =      int(lowerCamelCase_  )
														if self.entailment_id == -1:
																					raise ValueError('''Could not determine the entailment ID from the model config, please pass it at init.'''  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_  )			->				Any:
														lowerCAmelCase__			  =      labels
														return self.pre_processor(
														    [text] * len(lowerCamelCase_  )   ,       [F"""This example is {label}""" for label in labels]   ,       return_tensors='''pt'''   ,       padding='''max_length'''   ,       )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				str:
														lowerCAmelCase__			  =      outputs.logits
														lowerCAmelCase__			  =      torch.argmax(logits[:, 2]  ).item()
														return self._labels[label_id]



'''simple docstring'''






from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer




@dataclass
class 						a__							(				a__    ):





							'''simple docstring'''






							lowercase__     :  torch.FloatTensor




class 						a__							(				a__							,							a__    ):





							'''simple docstring'''






							@register_to_config
							def __init__(						self   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = ("DownEncoderBlock2D",)   ,       lowerCamelCase_ = ("UpDecoderBlock2D",)   ,       lowerCamelCase_ = (64,)   ,       lowerCamelCase_ = 1   ,       lowerCamelCase_ = "silu"   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = 32   ,       lowerCamelCase_ = 2_56   ,       lowerCamelCase_ = 32   ,       lowerCamelCase_ = None   ,       lowerCamelCase_ = 0.18_215   ,       lowerCamelCase_ = "group"   ,       )			->				Union[str, Any]:
														super().__init__()

														# pass init params to Encoder
														lowerCAmelCase__			  =      Encoder(
														    in_channels=lowerCamelCase_   ,       out_channels=lowerCamelCase_   ,       down_block_types=lowerCamelCase_   ,       block_out_channels=lowerCamelCase_   ,       layers_per_block=lowerCamelCase_   ,       act_fn=lowerCamelCase_   ,       norm_num_groups=lowerCamelCase_   ,       double_z=lowerCamelCase_   ,       )

														lowerCAmelCase__			  =      vq_embed_dim if vq_embed_dim is not None else latent_channels

														lowerCAmelCase__			  =      nn.Convad(lowerCamelCase_   ,       lowerCamelCase_   ,       1  )
														lowerCAmelCase__			  =      VectorQuantizer(lowerCamelCase_   ,       lowerCamelCase_   ,       beta=0.25   ,       remap=lowerCamelCase_   ,       sane_index_shape=lowerCamelCase_  )
														lowerCAmelCase__			  =      nn.Convad(lowerCamelCase_   ,       lowerCamelCase_   ,       1  )

														# pass init params to Decoder
														lowerCAmelCase__			  =      Decoder(
														    in_channels=lowerCamelCase_   ,       out_channels=lowerCamelCase_   ,       up_block_types=lowerCamelCase_   ,       block_out_channels=lowerCamelCase_   ,       layers_per_block=lowerCamelCase_   ,       act_fn=lowerCamelCase_   ,       norm_num_groups=lowerCamelCase_   ,       norm_type=lowerCamelCase_   ,       )







							@apply_forward_hook
							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = True  )			->				VQEncoderOutput:
														lowerCAmelCase__			  =      self.encoder(lowerCamelCase_  )
														lowerCAmelCase__			  =      self.quant_conv(lowerCamelCase_  )

														if not return_dict:
																					return (h,)

														return VQEncoderOutput(latents=lowerCamelCase_  )







							@apply_forward_hook
							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = False   ,       lowerCamelCase_ = True  )			->				Union[DecoderOutput, torch.FloatTensor]:
														# also go through quantization layer
														if not force_not_quantize:
																					lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__			  =      self.quantize(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      h
														lowerCAmelCase__			  =      self.post_quant_conv(lowerCamelCase_  )
														lowerCAmelCase__			  =      self.decoder(lowerCamelCase_   ,       quant if self.config.norm_type == '''spatial''' else None  )

														if not return_dict:
																					return (dec,)

														return DecoderOutput(sample=lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = True  )			->				Union[DecoderOutput, torch.FloatTensor]:
														lowerCAmelCase__			  =      sample
														lowerCAmelCase__			  =      self.encode(lowerCamelCase_  ).latents
														lowerCAmelCase__			  =      self.decode(lowerCamelCase_  ).sample

														if not return_dict:
																					return (dec,)

														return DecoderOutput(sample=lowerCamelCase_  )