Добавление нового конвейера

Сделайте Pipeline своим, создав подкласс и реализовав несколько методов. Поделитесь кодом с сообществом на Hub и зарегистрируйте пайплайн в Transformers, чтобы все могли быстро и легко им пользоваться.

В этом руководстве описан процесс добавления нового конвейера в Transformers.

Дизайнерские решения

Минимально достаточно передать в Pipeline подходящий вход для задачи. Именно с этого стоит начинать при проектировании конвейера.

Определите, какие типы входных данных может принимать Pipeline. Это могут быть строки, необработанные байты, словари и так далее. По возможности старайтесь использовать чистый Python для входных данных, так как это обеспечивает лучшую совместимость. Далее определите, какой вывод должен возвращать Pipeline. Снова самым простым и лучшим вариантом является использование Python для вывода, поскольку с ним проще работать.

Простые входные и выходные данные, в идеале сериализуемые в JSON, упрощают запуск вашего Pipeline для пользователей, избавляя их от необходимости изучать новые типы объектов. Также распространена практика поддержки множества различных типов входных данных для ещё большего удобства. Например, возможность принимать аудиофайл по имени файла, URL или в виде сырых байтов предоставляет пользователю больше гибкости в способе предоставления аудиоданных.

Создание конвейера

Когда входные и выходные данные определены, можно приступать к реализации Pipeline. Ваш пайплайн должен наследоваться от базового класса Pipeline и включать 4 метода.

from transformers import Pipeline

class MyPipeline(Pipeline):
    def _sanitize_parameters(self, **kwargs):

    def preprocess(self, inputs, args=2):

    def _forward(self, model_inputs):

    def postprocess(self, model_outputs):

preprocess принимает входные данные и преобразует их в подходящий для модели формат.

def preprocess(self, inputs, maybe_arg=2):
    model_input = Tensor(inputs["input_ids"])
    return {"model_input": model_input}

_forward не следует вызывать напрямую. forward — предпочтительный метод, поскольку он включает механизмы защиты, гарантирующие корректную работу на ожидаемом устройстве. Всё, что связано с моделью, должно находиться в _forward, а остальное — в preprocess или postprocess.

def _forward(self, model_inputs):
    outputs = self.model(**model_inputs)
    return outputs

postprocess формирует итоговый вывод на основе результата модели, полученного в _forward.

def postprocess(self, model_outputs, top_k=5):
    best_class = model_outputs["logits"].softmax(-1)
    return best_class

_sanitize_parameters позволяет пользователям передавать дополнительные параметры в Pipeline. Это может происходить при инициализации или при вызове Pipeline. _sanitize_parameters возвращает 3 словаря с дополнительными именованными аргументами, которые передаются непосредственно в preprocess, _forward и postprocess. Не добавляйте ничего, если пользователь не вызвал пайплайн с дополнительными параметрами. Это сохраняет аргументы по умолчанию в определении функции, что всегда более естественно.

Например, добавьте параметр top_k в postprocess, чтобы возвращать 5 наиболее вероятных классов. Затем в _sanitize_parameters проверьте, передал ли пользователь top_k, и добавьте его в postprocess_kwargs.

def _sanitize_parameters(self, **kwargs):
    preprocess_kwargs = {}
    if "maybe_arg" in kwargs:
        preprocess_kwargs["maybe_arg"] = kwargs["maybe_arg"]

    postprocess_kwargs = {}
    if "top_k" in kwargs:
        postprocess_kwargs["top_k"] = kwargs["top_k"]
    return preprocess_kwargs, {}, postprocess_kwargs

Теперь конвейер может возвращать наиболее вероятные метки, если пользователь этого пожелает.

from transformers import pipeline

pipeline = pipeline("my-task")
# returns 3 most likely labels
pipeline("This is the best meal I've ever had", top_k=3)
# returns 5 most likely labels by default
pipeline("This is the best meal I've ever had")

Регистрация конвейера

Зарегистрируйте новую задачу, которую поддерживает ваш конвейер, в PIPELINE_REGISTRY. Реестр определяет:

Поддерживаемый класс модели Pytorch с pt_model
модель по умолчанию, которая должна быть получена из конкретного ревизии (ветки или хэша коммита), где модель работает как ожидается с default
ожидаемый вход с type

from transformers.pipelines import PIPELINE_REGISTRY
from transformers import AutoModelForSequenceClassification

PIPELINE_REGISTRY.register_pipeline(
    "new-task",
    pipeline_class=MyPipeline,
    pt_model=AutoModelForSequenceClassification,
    default={"pt": ("user/awesome-model", "branch-name")},
    type="text",
)

Поделитесь своим конвейером

Поделитесь своим пайплайном с сообществом на Hub или добавьте его непосредственно в Transformers.

Загрузка кода вашего пайплайна на Hub происходит быстрее, так как не требует ревью от команды Transformers. Добавление пайплайна в Transformers может занять больше времени, поскольку требуется ревью, а также необходимо добавить тесты, чтобы убедиться, что ваш Pipeline работает.

Загрузка в Hub

Добавьте код вашего конвейера в Hub в файле 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}

Сохраните код в файле с именем pair_classification.py и импортируйте и зарегистрируйте его, как показано ниже.

from pair_classification import PairClassificationPipeline
from transformers.pipelines import PIPELINE_REGISTRY
from transformers import AutoModelForSequenceClassification

PIPELINE_REGISTRY.register_pipeline(
    "pair-classification",
    pipeline_class=PairClassificationPipeline,
    pt_model=AutoModelForSequenceClassification,
)

Функция register_pipeline регистрирует сведения о конвейере (тип задачи, класс конвейера, поддерживаемые бэкенды) в файле config.json модели.

  "custom_pipelines": {
    "pair-classification": {
      "impl": "pair_classification.PairClassificationPipeline",
      "pt": [
        "AutoModelForSequenceClassification"
      ],
    }
  },

Вызовите push_to_hub(), чтобы отправить пайплайн на Hub. Файл Python с кодом копируется на Hub, а модель и токенизатор пайплайна также сохраняются и отправляются на Hub. Теперь ваш пайплайн должен быть доступен на Hub в вашем пространстве имён.

from transformers import pipeline

pipeline = pipeline(task="pair-classification", model="sgugger/finetuned-bert-mrpc")
pipeline.push_to_hub("pair-classification-pipeline")

Чтобы использовать конвейер, добавьте trust_remote_code=True при загрузке конвейера.

from transformers import pipeline

pipeline = pipeline(task="pair-classification", trust_remote_code=True)

Добавить в Transformers

Добавление пользовательского конвейера в Transformers требует написания тестов для проверки корректности работы и запроса ревью от команды Transformers.

Добавьте код вашего конвейера в качестве нового модуля в подмодуль pipelines и включите его в список задач, определённых в pipelines/init.py.

Далее добавьте новый тест для пайплайна в transformers/tests/pipelines. Вы можете посмотреть другие тесты, чтобы увидеть примеры того, как тестировать ваш пайплайн.

Функция run_pipeline_test должна быть максимально универсальной и работать с моделями, определёнными в model_mapping. Это важно для проверки будущей совместимости с новыми моделями.

Вы также заметите, что ANY используется по всему коду функции run_pipeline_test. Модели случайные, поэтому проверить фактические значения невозможно. Использование ANY позволяет тесту сопоставлять вывод с типом пайплайна.

Наконец, вам также следует реализовать следующие 4 теста.

test_small_model_pt, используйте небольшую модель для этих конвейеров, чтобы убедиться, что они возвращают корректные результаты. Результаты не обязаны иметь смысл. Каждый конвейер должен возвращать одинаковый результат.
test_large_model_pt, используйте реалистичную модель для этих конвейеров, чтобы убедиться, что они возвращают осмысленные результаты. Эти тесты выполняются медленно и должны быть помечены как медленные.

Update on GitHub