Pipeline

Pipeline

이미 학습된 모델 하나만 있으면 어떤 작업이든 바로 추론해 보고 싶을 때가 있어요. Pipeline은 허깅페이스 허브의 어떤 모델로든 다양한 머신러닝 작업을 바로 실행할 수 있게 해 주는, 단순하지만 강력한 추론 API예요. 작업에 맞는 파라미터를 지정해 회의록 전사 같은 자동음성인식(ASR) 파이프라인에 타임스탬프를 붙이는 식으로 조정할 수도 있어요. GPU·애플 실리콘·반정밀도(half-precision) 가중치를 지원해서 추론을 가속하고 메모리도 아낄 수 있죠.

출처: Pipeline · Hugging Face Transformers

파이프라인 클래스

Transformers에는 범용 Pipeline 클래스와 TextGenerationPipeline 같은 여러 개별 작업별 파이프라인 클래스가 있어요. 개별 파이프라인은 pipeline의 task 파라미터에 작업 식별자를 지정해 불러올 수 있어요. 각 파이프라인의 작업 식별자는 해당 API 문서에서 찾을 수 있어요.

각 작업은 기본 프리트레인 모델과 전처리기(preprocessor)를 사용하도록 설정되어 있는데, 다른 모델을 쓰고 싶다면 model 파라미터로 덮어쓸 수 있어요.

예를 들어 Gemma 2로 텍스트 생성을 하려면 task="text-generation", model="google/gemma-2-2b"를 지정해요.

from transformers import pipeline

pipeline = pipeline(task="text-generation", model="google/gemma-2-2b")
pipeline("the secret to baking a really good cake is ")
[{'generated_text': 'the secret to baking a really good cake is 1. the right ingredients 2. the'}]

입력이 여러 개라면 리스트로 넘겨요. accelerate의 Accelerator로 디바이스도 정할 수 있어요.

from transformers import pipeline
from accelerate import Accelerator

device = Accelerator().device
pipeline = pipeline(task="text-generation", model="google/gemma-2-2b", device=device)

device 파라미터에 디바이스 인덱스를 직접 지정하는 방법도 있고, 허브의 어느 모델이든 파이프라인에 그대로 넘겨 쓸 수 있어요.

더 알아보기