Method sample

llms/speculative_decoding/decoder.py:110–111

(logits)

        max_tokens: int = 100,
    ):
        def sample(logits):
            return mx.argmax(logits, axis=-1)

        prompt = mx.array(self.tokenizer.encode(prompt), mx.uint32)[None]
        memory = self.model.encode(prompt)