Ajout d'un nouveau pipeline

Créez votre propre Pipeline en le sous-classant et en implémentant quelques méthodes. Partagez le code avec la communauté sur le Hub et enregistrez le pipeline dans Transformers pour que tout le monde puisse l’utiliser rapidement et facilement.

Ce guide vous accompagne dans le processus d’ajout d’un nouveau pipeline à Transformers.

Choix de conception

Au minimum, vous n'avez besoin de fournir à Pipeline qu'une entrée appropriée pour une tâche. C'est également le point de départ à privilégier lors de la conception de votre pipeline.

Définissez les types d’entrée que Pipeline peut accepter. Il peut s’agir de chaînes de caractères, d’octets bruts, de dictionnaires, etc. Essayez de conserver les entrées en Python pur lorsque c’est possible, car c’est plus compatible. Ensuite, déterminez la sortie que Pipeline doit retourner. Là encore, conserver la sortie en Python est l’option la plus simple et la meilleure, car elle est plus facile à manipuler.

Garder les entrées et les sorties simples, et idéalement sérialisables en JSON, facilite l’exécution de votre Pipeline par les utilisateurs, sans qu’ils aient à apprendre de nouveaux types d’objets. Il est également courant de prendre en charge de nombreux types d’entrées différents pour une facilité d’utilisation accrue. Par exemple, accepter un fichier audio via un nom de fichier, une URL ou des octets bruts offre à l’utilisateur plus de flexibilité dans la manière de fournir les données audio.

Créer un pipeline

Une fois l’entrée et la sortie définies, vous pouvez commencer à implémenter Pipeline. Votre pipeline doit hériter de la classe de base Pipeline et inclure 4 méthodes.

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 prend les entrées et les transforme dans le format d'entrée approprié pour le modèle.

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

_forward ne doit pas être appelé directement. forward est la méthode recommandée car elle inclut des garde-fous pour s’assurer que tout fonctionne correctement sur l’appareil attendu. Tout ce qui est lié au modèle doit être placé dans _forward et tout le reste dans preprocess ou postprocess.

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

postprocess génère la sortie finale à partir de la sortie du modèle dans _forward.

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

_sanitize_parameters permet aux utilisateurs de passer des paramètres supplémentaires à Pipeline. Cela peut se faire lors de l’initialisation ou lors de l’appel de Pipeline. _sanitize_parameters renvoie 3 dictionnaires d’arguments supplémentaires par mot-clé qui sont transmis directement à preprocess, _forward et postprocess. N’ajoutez rien si l’utilisateur n’a pas appelé le pipeline avec des paramètres supplémentaires. Cela conserve les arguments par défaut dans la définition de la fonction, ce qui est toujours plus naturel.

Par exemple, ajoutez un paramètre top_k dans postprocess pour retourner les 5 classes les plus probables. Ensuite, dans _sanitize_parameters, vérifiez si l'utilisateur a passé top_k et ajoutez-le à 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

Le pipeline peut désormais retourner les étiquettes les plus probables si l'utilisateur le souhaite.

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

Enregistrer un pipeline

Enregistrez la nouvelle tâche prise en charge par votre pipeline dans le PIPELINE_REGISTRY. Le registre définit :

La classe de modèle Pytorch prise en charge avec pt_model
un modèle par défaut qui doit provenir d'une révision spécifique (branche ou hachage de commit) où le modèle fonctionne comme prévu avec default
l'entrée attendue avec 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",
)

Partagez votre pipeline

Partagez votre pipeline avec la communauté sur le Hub ou vous pouvez l’ajouter directement à Transformers.

Il est plus rapide de téléverser votre code de pipeline sur le Hub car cela ne nécessite pas de relecture par l’équipe Transformers. L’ajout du pipeline à Transformers peut être plus lent car il exige une relecture et vous devez ajouter des tests pour vous assurer que votre Pipeline fonctionne.

Téléverser sur le Hub

Ajoutez le code de votre pipeline au Hub dans un fichier Python.

Par exemple, un pipeline personnalisé pour la classification de paires de phrases pourrait ressembler au code ci-dessous.

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}

Enregistrez le code dans un fichier nommé pair_classification.py, puis importez-le et enregistrez-le comme indiqué ci-dessous.

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 fonction register_pipeline enregistre les détails du pipeline (type de tâche, classe de pipeline, backends pris en charge) dans le fichier config.json d'un modèle.

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

Appel de push_to_hub() pour pousser le pipeline vers le Hub. Le fichier Python contenant le code est copié sur le Hub, et le modèle et le tokenizer du pipeline sont également enregistrés et poussés vers le Hub. Votre pipeline devrait maintenant être disponible sur le Hub sous votre espace de noms.

from transformers import pipeline

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

Pour utiliser le pipeline, ajoutez trust_remote_code=True lors du chargement du pipeline.

from transformers import pipeline

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

Ajouter à Transformers

Ajouter un pipeline personnalisé à Transformers nécessite d’ajouter des tests pour s’assurer que tout fonctionne comme prévu, ainsi que de demander une revue à l’équipe Transformers.

Ajoutez votre code de pipeline en tant que nouveau module au sous-module pipelines, et ajoutez-le à la liste des tâches définies dans pipelines/init.py.

Ensuite, ajoutez un nouveau test pour le pipeline dans transformers/tests/pipelines. Vous pouvez consulter les autres tests pour des exemples de la manière de tester votre pipeline.

La fonction run_pipeline_test doit être très générique et s’exécuter sur les modèles définis dans model_mapping. C’est important pour tester la compatibilité future avec de nouveaux modèles.

Vous remarquerez également que ANY est utilisé tout au long de la fonction run_pipeline_test. Les modèles étant aléatoires, il est impossible de vérifier les valeurs réelles. L’utilisation de ANY permet au test de correspondre au type de sortie du pipeline.

Enfin, vous devez également implémenter les 4 tests suivants.

test_small_model_pt, utilisez un petit modèle pour ces pipelines afin de vous assurer qu’ils renvoient les sorties correctes. Les résultats n’ont pas besoin d’avoir de sens. Chaque pipeline doit renvoyer le même résultat.
test_large_model_pt, utilisez un modèle réaliste pour ces pipelines afin de vous assurer qu'ils renvoient des résultats significatifs. Ces tests sont lents et doivent être marqués comme lents.

Update on GitHub