INNER CODE UNIT · Python

StreamModel

hyperonym/basaran · basaran/model.py:21

class StreamModel:
    """StreamModel wraps around a language model to provide stream decoding."""

    def __init__(self, model, tokenizer):
        super().__init__()
        self.tokenizer = tokenizer
        if torch.backends.mps.is_available():
            self.device = "mps"
        elif torch.cuda.is_available():
            self.device = "cuda"
        else:
            self.device = "cpu"
        self.model = model.to(self.device)

    def __call__(
        self,
        prompt,
        min_tokens=0,

View source record →