|
import numpy as np |
|
|
|
from transformers import Pipeline |
|
from ultravox.model import ultravox_model |
|
|
|
|
|
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} |
|
|