| 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} |
|
|
