File size: 803 Bytes
692d9ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from transformers import Pipeline
from .classifier import Classifier
class MultiTaskClassifierPipeline(Pipeline):
def _sanitize_parameters(self, **kwargs):
preprocess_kwargs = {}
postprocess_kwargs = {}
return preprocess_kwargs, {}, postprocess_kwargs
def preprocess(self, inputs):
return self.tokenizer(inputs, padding="max_length", truncation=True, return_tensors=self.framework).to(self.device)
def _forward(self, model_inputs):
return self.model(**model_inputs)
def postprocess(self, model_outputs):
model_config = self.model.config
classifier = Classifier(model_config.task_specific_params[Classifier.MODEL_CONFIG])
logits = model_outputs.logits.numpy()
return classifier.get_results(logits)[0] |