






import torch

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



class  __lowerCAmelCase     (	a__   ):
 lowerCAmelCase__       					=  """facebook/bart-large-mnli"""
 lowerCAmelCase__       					=  (
     """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."""
 )
 lowerCAmelCase__       					=  """text_classifier"""
 lowerCAmelCase__       					=  AutoTokenizer
 lowerCAmelCase__       					=  AutoModelForSequenceClassification

 lowerCAmelCase__       					=  ["""text""", ["""text"""]]
 lowerCAmelCase__       					=  ["""text"""]




 def        lowerCamelCase      (						self				):



    '''simple docstring'''
    super().setup()
    __lowerCamelCase													=    self.model.config
    __lowerCamelCase													=    -1
    for idx, label in config.idalabel.items():
       if label.lower().startswith('''entail'''				):
          __lowerCamelCase													=    int(lowercase__				)
    if self.entailment_id == -1:
       raise ValueError('''Could not determine the entailment ID from the model config, please pass it at init.'''				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    labels
    return self.pre_processor(
        [text] * len(lowercase__				)	,				[F"""This example is {label}""" for label in labels]	,				return_tensors='''pt'''	,				padding='''max_length'''	,				)




 def        lowerCamelCase      (						self	,				__UpperCAmelCase				):



    '''simple docstring'''
    __lowerCamelCase													=    outputs.logits
    __lowerCamelCase													=    torch.argmax(logits[:, 2]				).item()
    return self._labels[label_id]











import torch

from diffusers import StableDiffusionPipeline


a_			=					"""path-to-your-trained-model"""
a_			=					StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.floataa).to("""cuda""")

a_			=					"""A photo of sks dog in a bucket"""
a_			=					pipe(prompt, num_inference_steps=50, guidance_scale=7.5).images[0]

image.save("""dog-bucket.png""")






