HCKLab/BiBert-MultiTask-1

1from transformers import Pipeline2import numpy as np3import torch4 5def softmax(_outputs):6  maxes = np.max(_outputs, axis=-1, keepdims=True)7  shifted_exp = np.exp(_outputs - maxes)8  return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True)9 10class BiBert_MultiTaskPipeline(Pipeline):11 12  def _sanitize_parameters(self, **kwargs):13 14    preprocess_kwargs = {}15    if "task_id" in kwargs:16      preprocess_kwargs["task_id"] = kwargs["task_id"]17    18    forward_kwargs = {}19    if "task_id" in kwargs:20      forward_kwargs["task_id"] = kwargs["task_id"]21    return preprocess_kwargs, forward_kwargs, {}22 23  def preprocess(self, inputs, task_id):24    return_tensors = self.framework25    feature = self.tokenizer(inputs, padding = True, return_tensors=return_tensors).to(self.device)26    task_ids = np.full(shape=1,fill_value=task_id, dtype=int)27    feature["task_ids"] = torch.IntTensor(task_ids)28    return feature29  30  def _forward(self, model_inputs, task_id):31    return self.model(**model_inputs)32 33  def postprocess(self, model_outputs):34    outputs = model_outputs["logits"][0]35    outputs = outputs.numpy()36    scores = softmax(outputs)37 38    dict_scores = [39            {"label": self.model.config.id2label[i], "score": score.item()} for i, score in enumerate(scores)40        ]41    return dict_scores42