Agregar un nuevo pipeline

Haz Pipeline a tu medida mediante la herencia y la implementación de algunos métodos. Comparte el código con la comunidad en el Hub y registra el pipeline en Transformers para que todos puedan usarlo de forma rápida y sencilla.

Esta guía te guiará a través del proceso de añadir un nuevo pipeline a Transformers.

Decisiones de diseño

Como mínimo, solo necesitas proporcionar a Pipeline una entrada adecuada para una tarea. Este es también el punto de partida al diseñar tu pipeline.

Decide qué tipos de entrada puede aceptar Pipeline. Pueden ser cadenas, bytes crudos, diccionarios, etc. Intenta mantener las entradas en Python puro siempre que sea posible, ya que es más compatible. A continuación, decide qué salida debe devolver Pipeline. De nuevo, mantener la salida en Python es la opción más simple y mejor, ya que es más fácil de trabajar.

Mantener las entradas y salidas simples, e idealmente serializables en JSON, facilita que los usuarios ejecuten su Pipeline sin necesidad de aprender nuevos tipos de objetos. También es habitual admitir muchos tipos de entrada diferentes para una mayor facilidad de uso. Por ejemplo, aceptar un archivo de audio desde un nombre de archivo, una URL o bytes en bruto otorga al usuario más flexibilidad en cómo proporciona los datos de audio.

Crear un pipeline

Con la entrada y la salida decididas, puedes comenzar a implementar Pipeline. Tu pipeline debe heredar de la clase 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):

1. preprocess toma las entradas y las transforma al formato de entrada adecuado para el modelo.

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

2. _forward no debe llamarse directamente. forward es el método preferido porque incluye salvaguardas para asegurarse de que todo funcione correctamente en el dispositivo esperado. Todo lo vinculado al modelo debe ir en _forward y todo lo demás en preprocess o postprocess.

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

3. postprocess genera la salida final a partir de la salida del modelo en _forward.

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

4. _sanitize_parameters permite a los usuarios pasar parámetros adicionales a Pipeline. Esto puede ocurrir durante la inicialización o cuando se llama a Pipeline. _sanitize_parameters devuelve 3 diccionarios de argumentos de palabra clave adicionales que se pasan directamente a preprocess, _forward y postprocess. No añada nada si el usuario no llamó a la pipeline con parámetros extra. Esto mantiene los argumentos predeterminados en la definición de la función, lo cual es siempre más natural.

Por ejemplo, añade un parámetro top_k en postprocess para devolver las 5 clases más probables. Luego, en _sanitize_parameters, comprueba si el usuario ha pasado top_k y añádelo 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

Ahora, el pipeline puede devolver las etiquetas más probables si el usuario lo elige.

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 un pipeline

Registra la nueva tarea que tu pipeline admite en PIPELINE_REGISTRY. El registro define:

• La clase de modelo de Pytorch compatible con pt_model
• un modelo predeterminado que debe provenir de una revisión específica (rama o hash de commit) donde el modelo funciona como se espera con default
• la entrada esperada con 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",
)

Comparte tu pipeline

Comparte tu pipeline con la comunidad en el Hub o puedes añadirlo directamente a Transformers.

Es más rápido subir el código de tu pipeline al Hub porque no requiere una revisión del equipo de Transformers. Añadir el pipeline a Transformers puede ser más lento porque requiere una revisión y debes añadir pruebas para asegurarte de que tu Pipeline funciona.

Subir al Hub

Agrega el código de tu pipeline al Hub en un archivo de Python.

Por ejemplo, un pipeline personalizado para la clasificación de pares de oraciones podría parecerse al siguiente código.

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}

Guarde el código en un archivo llamado pair_classification.py e impórtelo y regístrelo como se muestra a continuación.

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

La función register_pipeline registra los detalles del pipeline (tipo de tarea, clase de pipeline, backends compatibles) en el archivo config.json de los modelos.

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

Llama a push_to_hub() para subir el pipeline al Hub. El archivo Python que contiene el código se copia al Hub, y el modelo y el tokenizador del pipeline también se guardan y se suben al Hub. Tu pipeline debería estar disponible en el Hub bajo tu espacio de nombres.

from transformers import pipeline

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

Para usar el pipeline, agregue trust_remote_code=True al cargar el pipeline.

from transformers import pipeline

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

Agregar a Transformers

Para añadir un pipeline personalizado a Transformers, es necesario agregar pruebas para asegurarse de que todo funcione como se espera y solicitar una revisión al equipo de Transformers.

Agrega el código de tu pipeline como un nuevo módulo al submódulo pipelines y añádelo a la lista de tareas definidas en pipelines/init.py.

A continuación, añade una nueva prueba para el pipeline en transformers/tests/pipelines. Puedes consultar las otras pruebas como ejemplos de cómo probar tu pipeline.

La función run_pipeline_test debe ser muy genérica y ejecutarse en los modelos definidos en model_mapping. Esto es importante para probar la compatibilidad futura con nuevos modelos.

También notará que ANY se utiliza a lo largo de la función run_pipeline_test. Los modelos son aleatorios, por lo que no se pueden verificar los valores reales. El uso de ANY permite que la prueba coincida con la salida del tipo de pipeline.

Por último, también debes implementar las siguientes 4 pruebas.

1. test_small_model_pt, usa un modelo pequeño para estos pipelines para asegurarte de que devuelvan las salidas correctas. Los resultados no tienen que tener sentido. Cada pipeline debe devolver el mismo resultado.
2. test_large_model_pt, usa un modelo realista para estos pipelines para asegurarte de que devuelvan resultados significativos. Estas pruebas son lentas y deberían marcarse como lentas.