Predictor로 예측 정의하기

Cog에서 모델의 예측 로직은 Predictor 클래스로 정의해요. setup()으로 모델을 초기화하고, predict()로 예측을 수행하는 구조예요. 이 클래스가 곧 모델의 입력·출력 스키마를 결정해요.

출처: https://github.com/replicate/cog

기본 Predictor

BasePredictor를 상속받고 predict 메서드를 구현해요. 입력 파라미터의 타입 힌트가 그대로 모델의 입력 스키마가 되고, 반환 타입이 출력 스키마가 돼요.

from cog import BasePredictor, Input, Path

class Predictor(BasePredictor):
    def setup(self):
        """컨테이너 시작 시 한 번 모델을 메모리에 로드해요."""
        self.model = load_model("./weights")

    def predict(self, prompt: str, steps: int = 50) -> Path:
        """예측 요청마다 한 번씩 실행돼요."""
        output = self.model.generate(prompt, steps=steps)
        output.save("/tmp/output.png")
        return Path("/tmp/output.png")

setup(): 컨테이너 시작 시 한 번 호출돼요. 모델을 메모리에 올리는 무거운 작업을 여기서 해요.
predict(): 예측 요청마다 호출돼요. 시그니처가 입력 스키마를, 반환 타입이 출력 스키마를 결정해요. 동기(def)나 비동기(async def) 모두 가능해요.
train()(선택): 파인튜닝 워크플로를 위한 메서드로, predict와 같은 계약을 가져요. cog.yaml의 train 키로 별도 설정해요.

여러 출력과 스트리밍

여러 개를 반환할 수 있고, 결과를 점진적으로 흘려보낼 수도 있어요.

여러 출력은 리스트로 반환하고요.

from typing import List
from cog import Path

def predict(self, prompt: str) -> List[Path]:
    paths = []
    for i in range(4):
        path = f"/tmp/output_{i}.png"
        self.model.generate(prompt, seed=i).save(path)
        paths.append(Path(path))
    return paths

스트리밍은 Iterator로 값을 yield하면 돼요.

from typing import Iterator

def predict(self, prompt: str) -> Iterator[str]:
    for token in self.model.generate_stream(prompt):
        yield token

스트리밍 스키마는 x-cog-array-type: iterator로 표시돼요. 토큰 단위로 생성되는 결과를 실시간으로 받고 싶을 때 유용해요.

더 알아보기

환경 정의는 «cog.yaml 이해하기»를 보세요.
모델을 실제로 만들어 보려면 «첫 모델 만들기»를 확인하세요.