| from transformers import Pipeline | |
| import numpy as np | |
| import torch | |
| def softmax(_outputs): | |
| maxes = np.max(_outputs, axis=-1, keepdims=True) | |
| shifted_exp = np.exp(_outputs - maxes) | |
| return shifted_exp / shifted_exp.sum(axis=-1, keepdims=True) | |
| class BiBert_MultiTaskPipeline(Pipeline): | |
| def _sanitize_parameters(self, **kwargs): | |
| preprocess_kwargs = {} | |
| if "task_id" in kwargs: | |
| preprocess_kwargs["task_id"] = kwargs["task_id"] | |
| forward_kwargs = {} | |
| if "task_id" in kwargs: | |
| forward_kwargs["task_id"] = kwargs["task_id"] | |
| postprocess_kwargs = {} | |
| if "top_k" in kwargs: | |
| postprocess_kwargs["top_k"] = kwargs["top_k"] | |
| postprocess_kwargs["_legacy"] = False | |
| return preprocess_kwargs, forward_kwargs, postprocess_kwargs | |
| def preprocess(self, inputs, task_id): | |
| return_tensors = self.framework | |
| feature = self.tokenizer(inputs, padding = True, return_tensors=return_tensors).to(self.device) | |
| task_ids = np.full(shape=1,fill_value=task_id, dtype=int) | |
| feature["task_ids"] = torch.IntTensor(task_ids) | |
| return feature | |
| def _forward(self, model_inputs, task_id): | |
| return self.model(**model_inputs) | |
| def postprocess(self, model_outputs, top_k=1, _legacy=True): | |
| outputs = model_outputs["logits"][0] | |
| outputs = outputs.numpy() | |
| scores = softmax(outputs) | |
| if top_k == 1 and _legacy: | |
| return {"label": self.model.config.id2label[scores.argmax().item()], "score": scores.max().item()} | |
| dict_scores = [ | |
| {"label": self.model.config.id2label[i], "score": score.item()} for i, score in enumerate(scores) | |
| ] | |
| if not _legacy: | |
| dict_scores.sort(key=lambda x: x["score"], reverse=True) | |
| if top_k is not None: | |
| dict_scores = dict_scores[:top_k] | |
| return dict_scores | |