cog
cog copied to clipboard
A support for `BaseModel` subclass output streaming
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]:
...