File size: 1,422 Bytes
6511d8d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
print("Loading Multi head pipeline")
from transformers.pipelines import PIPELINE_REGISTRY
from transformers import TextClassificationPipeline, AutoTokenizer, AutoModelForSequenceClassification
class CustomTextClassificationPipeline(TextClassificationPipeline):
def __init__(self, model, tokenizer=None, **kwargs):
if tokenizer is None:
tokenizer = AutoTokenizer.from_pretrained(model.config._name_or_path)
super().__init__(model=model, tokenizer=tokenizer, **kwargs)
def _sanitize_parameters(self, **kwargs):
preprocess_kwargs = {}
return preprocess_kwargs, {}, {}
def preprocess(self, inputs):
return self.tokenizer(inputs, return_tensors='pt', truncation=False)
def _forward(self, model_inputs):
input_ids = model_inputs['input_ids']
attention_mask = (input_ids != 0).long()
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
return outputs
def postprocess(self, model_outputs):
predictions = model_outputs.logits.argmax(dim=-1).squeeze().tolist()
categories = ["Race/Origin", "Gender/Sex", "Religion", "Ability", "Violence", "Other"]
return dict(zip(categories, predictions))
PIPELINE_REGISTRY.register_pipeline(
"multi-head-text-classification",
pipeline_class=CustomTextClassificationPipeline,
pt_model=AutoModelForSequenceClassification,
) |