새 파이프라인 추가하기

새 파이프라인 추가하기 (Adding a new pipeline)

Pipeline을 서브클래싱하고 몇 가지 메서드를 구현해서 나만의 것으로 만들 수 있어요. 커뮤니티와 코드를 Hub에서 공유하고, Transformers에 파이프라인을 등록하면 누구나 빠르고 쉽게 사용할 수 있지요. 이 가이드에서는 Transformers에 새 파이프라인을 추가하는 전체 과정을 살펴볼게요.

출처: 문서

본문

설계 선택 (Design choices)

최소한 Pipeline에는 태스크에 적합한 입력만 제공하면 돼요. 파이프라인을 설계할 때도 여기서 시작하는 게 좋아요.

Pipeline이 받을 수 있는 입력 타입을 결정해요. 문자열, 원시 바이트(raw bytes), 딕셔너리 등이 될 수 있어요. 가능하면 입력을 순수 Python으로 유지하는 게 좋아요. 호환성이 더 좋기 때문이지요. 다음으로 Pipeline이 반환해야 할 출력을 결정해요. 역시 출력을 Python으로 유지하는 게 가장 단순하고 좋은 선택이에요. 다루기 더 쉽기 때문이에요.

입력과 출력을 단순하게, 이상적으로는 JSON 직렬화 가능하게 유지하면 사용자가 새 객체 타입을 배울 필요 없이 여러분의 Pipeline을 실행할 수 있어요. 또한 더 큰 사용 편의를 위해 다양한 입력 타입을 지원하는 것도 흔해요. 예를 들어 오디오 파일을 파일 이름, URL, 원시 바이트 중 어떤 것으로도 받을 수 있게 하면 사용자가 오디오 데이터를 제공하는 방식에 더 유연해지지요.

파이프라인 만들기 (Create a pipeline)

입력과 출력을 정했다면 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는 preprocess, _forward, postprocess에 직접 전달되는 추가 키워드 인자 3개의 dict를 반환해요. 사용자가 추가 파라미터 없이 파이프라인을 호출했다면 아무것도 추가하지 마세요. 그러면 함수 정의에 있는 기본 인자가 그대로 유지되는데, 이게 항상 더 자연스러워요.

예를 들어 postprocess에 top_k 파라미터를 추가해서 가장 가능성 높은 상위 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")

파이프라인 등록하기 (Register a pipeline)

파이프라인이 지원하는 새 태스크를 PIPELINE_REGISTRY에 등록해요. 레지스트리는 다음을 정의해요:

pt_model로 지원되는 Pytorch 모델 클래스
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에 업로드하기

파이프라인 코드를 Python 파일로 Hub에 추가해요.

예를 들어 문장 쌍 분류(sentence pair classification)를 위한 커스텀 파이프라인은 아래 코드처럼 생겼을 거예요.

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라는 파일에 저장하고, 아래처럼 import 해서 등록해요.

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에 복사되고, 파이프라인의 모델과 tokenizer도 저장되어 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 함수는 아주 일반적(generic)이어야 하고 model_mapping에 정의된 모델들에서 실행되어야 해요. 이는 새 모델과의 향후 호환성을 테스트하는 데 중요해요.

run_pipeline_test 함수 전반에 ANY가 쓰이는 것도 알 수 있을 거예요. 모델이 랜덤이라 실제 값을 확인할 수 없기 때문이에요. ANY를 사용하면 테스트가 대신 파이프라인 타입의 출력과 일치하도록 해줘요.

마지막으로 다음 2가지 테스트도 구현해야 해요.

test_small_model_pt — 이 파이프라인들에 작은 모델을 사용해서 올바른 출력을 반환하는지 확인해요. 결과가 말이 될 필요는 없어요. 각 파이프라인이 같은 결과를 반환하면 돼요.
test_large_model_pt — 이 파이프라인들에 현실적인 모델을 사용해서 의미 있는 결과를 반환하는지 확인해요. 이 테스트들은 느리므로 slow로 표시해야 해요.

더 알아보기 (Learn more)

Pipeline API

새 파이프라인 추가하기