Method _generate

llms/speculative_decoding/decoder.py:55–67

(
        self,
        x: mx.array,
        memory: mx.array,
        draft: bool = False,
    )

        self.delta = delta

    def _generate(
        self,
        x: mx.array,
        memory: mx.array,
        draft: bool = False,
    ):
        model = self.draft_model if draft else self.model
        while True:
            logits = model.decode(x[None], memory)[0, -1]
            x = mx.argmax(logits, keepdims=True)
            lognorm = mx.logsumexp(logits.astype(mx.float32))
            logprob = logits[x] - lognorm
            yield x, logprob

    def generate(
        self,