34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
import numpy as np
|
|
|
|
from transformers import Pipeline
|
|
|
|
|
|
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 PairClassificationPipeline(Pipeline):
|
|
def _sanitize_parameters(self, **kwargs):
|
|
preprocess_kwargs = {}
|
|
if "second_text" in kwargs:
|
|
preprocess_kwargs["second_text"] = kwargs["second_text"]
|
|
return preprocess_kwargs, {}, {}
|
|
|
|
def preprocess(self, text, second_text=None):
|
|
return self.tokenizer(text, text_pair=second_text, return_tensors=self.framework)
|
|
|
|
def _forward(self, model_inputs):
|
|
return self.model(**model_inputs)
|
|
|
|
def postprocess(self, model_outputs):
|
|
logits = model_outputs.logits[0].numpy()
|
|
probabilities = softmax(logits)
|
|
|
|
best_class = np.argmax(probabilities)
|
|
label = self.model.config.id2label[best_class]
|
|
score = probabilities[best_class].item()
|
|
logits = logits.tolist()
|
|
return {"label": label, "score": score, "logits": logits}
|