cog icon indicating copy to clipboard operation
cog copied to clipboard

A support for `BaseModel` subclass output streaming

Open sitatec opened this issue 1 year ago • 0 comments

Currently only str, int, float, bool, or cog.Path are supported for streaming. It would be useful to support custom objects extending BaseModel, which can be used to include some metadata:

from cog import BasePredictor, BaseModel
from typing import Iterator

class Output(BaseModel):
    token_count: int
    safety_level: str
    ...
    text: str

class Predictor(BasePredictor):
    def predict(self) -> Iterator[Output]:
          ...

sitatec avatar Mar 26 '25 17:03 sitatec