Adicionando um novo pipeline

Crie o seu próprio Pipeline herdando a classe e implementando alguns métodos. Compartilhe o código com a comunidade no Hub e registre o pipeline no Transformers para que todos possam usá-lo de forma rápida e fácil.

Este guia mostrará como adicionar um novo pipeline ao Transformers.

Escolhas de design

No mínimo, você só precisa fornecer ao Pipeline uma entrada adequada para uma tarefa. Este também é o ponto de partida ao projetar seu pipeline.

Defina os tipos de entrada que o Pipeline pode aceitar. Podem ser strings, bytes brutos, dicionários, entre outros. Sempre que possível, mantenha as entradas em Python puro, pois isso oferece maior compatibilidade. Em seguida, defina a saída que o Pipeline deve retornar. Novamente, manter a saída em Python é a opção mais simples e recomendada, pois facilita o trabalho.

Manter as entradas e saídas simples e, idealmente, serializáveis em JSON, facilita para os usuários executarem seu Pipeline sem a necessidade de aprender novos tipos de objetos. Também é comum suportar muitos tipos diferentes de entrada para uma usabilidade ainda maior. Por exemplo, aceitar um arquivo de áudio a partir de um nome de arquivo, URL ou bytes brutos dá ao usuário mais flexibilidade na forma de fornecer os dados de áudio.

Criar um pipeline

Com a entrada e a saída definidas, você pode começar a implementar o Pipeline. Seu pipeline deve herdar da classe base Pipeline e incluir 4 métodos.

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 recebe as entradas e as transforma no formato de entrada adequado para o modelo.

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

_forward não deve ser chamado diretamente. forward é o método preferido, pois inclui salvaguardas para garantir que tudo funcione corretamente no dispositivo esperado. Tudo que estiver vinculado ao modelo deve ficar em _forward, e o restante deve ficar em preprocess ou postprocess.

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

postprocess gera a saída final a partir da saída do modelo em _forward.

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

_sanitize_parameters permite que os usuários passem parâmetros adicionais para Pipeline. Isso pode ocorrer durante a inicialização ou quando Pipeline é chamado. _sanitize_parameters retorna 3 dicts de argumentos de palavra-chave adicionais que são passados diretamente para preprocess, _forward e postprocess. Não adicione nada se o usuário não tiver chamado o pipeline com parâmetros extras. Isso preserva os argumentos padrão na definição da função, o que é sempre mais natural.

Por exemplo, adicione um parâmetro top_k em postprocess para retornar as 5 classes mais prováveis. Em seguida, em _sanitize_parameters, verifique se o usuário passou top_k e adicione-o a 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

Agora, o pipeline pode retornar os rótulos mais prováveis, caso o usuário assim deseje.

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")

Registrar um pipeline

Registre a nova tarefa que seu pipeline suporta no PIPELINE_REGISTRY. O registro define:

A classe de modelo Pytorch suportada com pt_model
um modelo padrão que deve vir de uma revisão específica (branch ou hash de commit) em que o modelo funciona como esperado com default
a entrada esperada com 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",
)

Compartilhe seu pipeline

Compartilhe seu pipeline com a comunidade no Hub ou você pode adicioná-lo diretamente ao Transformers.

É mais rápido enviar o código do seu pipeline para o Hub, pois não exige uma revisão da equipe do Transformers. Adicionar o pipeline ao Transformers pode ser mais lento, pois requer uma revisão e você precisa adicionar testes para garantir que seu Pipeline funciona.

Enviar para o Hub

Adicione o código do seu pipeline ao Hub em um arquivo Python.

Por exemplo, um pipeline personalizado para classificação de pares de frases pode ser semelhante ao código a seguir.

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}

Salve o código em um arquivo chamado pair_classification.py e importe e registre-o conforme mostrado abaixo.

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,
)

A função register_pipeline registra os detalhes do pipeline (tipo de tarefa, classe do pipeline, backends suportados) no arquivo config.json de um modelo.

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

Chame push_to_hub() para enviar o pipeline para o Hub. O arquivo Python contendo o código é copiado para o Hub, e o modelo e o tokenizer do pipeline também são salvos e enviados para o Hub. Seu pipeline deve estar disponível no Hub sob seu namespace.

from transformers import pipeline

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

Para usar o pipeline, adicione trust_remote_code=True ao carregar o pipeline.

from transformers import pipeline

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

Adicionar ao Transformers

Adicionar um pipeline personalizado ao Transformers exige a inclusão de testes para garantir que tudo funcione conforme o esperado, bem como a solicitação de uma revisão da equipe do Transformers.

Adicione o código do seu pipeline como um novo módulo ao submódulo pipelines e inclua-o na lista de tarefas definida em pipelines/init.py.

Em seguida, adicione um novo teste para o pipeline em transformers/tests/pipelines. Você pode consultar os outros testes como exemplos de como testar seu pipeline.

A função run_pipeline_test deve ser muito genérica e executar nos modelos definidos em model_mapping. Isso é importante para testar a compatibilidade futura com novos modelos.

Você também notará que ANY é usado ao longo da função run_pipeline_test. Os modelos são aleatórios, então não é possível verificar os valores reais. O uso de ANY permite que o teste corresponda ao tipo de saída do pipeline.

Por fim, você também deve implementar os 4 testes a seguir.

test_small_model_pt, use um modelo pequeno para esses pipelines para garantir que eles retornem as saídas corretas. Os resultados não precisam fazer sentido. Cada pipeline deve retornar o mesmo resultado.
test_large_model_pt, use um modelo realista para esses pipelines para garantir que retornem resultados significativos. Esses testes são lentos e devem ser marcados como lentos.

Update on GitHub