INNER CODE UNIT · Python

__init__

hyperonym/basaran · basaran/model.py:24

    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,
        max_tokens=16,
        temperature=1.0,
        top_p=1.0,

View source record →