zeroshot.py

from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.modeling_utils import PreTrainedModel
from transformers.configuration_utils import PretrainedConfig
import torch
from transformers import ZeroShotClassificationPipeline


class CustomConfig(PretrainedConfig):
 model_type = "test-zeroshot"

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


class CustomModel(PreTrainedModel):
 config_class = CustomConfig

 def __init__(self, config: CustomConfig):
 super().__init__(config)
 self.config = config
 self.embeddings = torch.nn.Embedding(num_embeddings=1, embedding_dim=1)

 def forward(self, **kwargs) -> SequenceClassifierOutput:

 return SequenceClassifierOutput(logits=torch.tensor([[1, 2, 3]]))


from transformers.pipelines import PIPELINE_REGISTRY

from transformers import AutoModelForSequenceClassification, TFAutoModelForSequenceClassification

if __name__ == "__main__":
 from transformers import pipeline
 classifier = pipeline("zero-shot-classification",
 model="cl-tohoku/bert-base-japanese")
 from transformers import AutoConfig, AutoModel, AutoModelForImageClassification

 CustomConfig.register_for_auto_class()
 CustomModel.register_for_auto_class("AutoModel")

 p = ZeroShotClassificationPipeline(model=CustomModel(CustomConfig()),
 tokenizer=classifier.tokenizer)
 from huggingface_hub import Repository

 repo = Repository("zero-shot-classification",
 clone_from="paulhindemith/zero-shot-classification")
 p.save_pretrained("zero-shot-classification")
 repo.push_to_hub()