/
githubmirror
/
transformers
Обзор
Документация
Войти
/
githubmirror
/
transformers
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
utils/test_module/custom_pipeline.py
33 строки
1 KB
Cyril Vallez
🚨🚨🚨 Fully remove Tensorflow and Jax support library-wide (#40760)
18 сен 2025, 19:27
Не верифицирован
18 сен 2025, 19:27
4df2529
Код
Авторство
О чём код?
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="pt") 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}