mlx-examples/phi2/model.py

from typing import Optional
from dataclasses import dataclass
from mlx.utils import tree_unflatten, tree_map
from transformers import AutoTokenizer

import mlx.core as mx
import mlx.nn as nn
import math


@dataclass
class ModelArgs:
    max_sequence_length: int = 2048
    num_vocab: int = 51200
    model_dim: int = 2560
    num_heads: int = 32
    num_layers: int = 32
    rotary_dim: int = 32


class NewGELUActivation(nn.Module):
    """
    Implementation of the GELU activation function currently in Google BERT repo (identical to OpenAI GPT). Also see
    the Gaussian Error Linear Units paper: https://arxiv.org/abs/1606.08415
    """

    def __call__(self, input: mx.array) -> mx.array:
        return (
            0.5
            * input
            * (
                1.0
                + mx.tanh(math.sqrt(2.0 / math.pi) * (input + 0.044715 * (input**3)))
            )
        )


class RoPEAttention(nn.Module):
    def __init__(self, dims: int, num_heads: int, bias: bool = True):
        super().__init__()

        self.num_heads = num_heads

        self.rope = nn.RoPE(dims // num_heads, traditional=True)
        self.query_proj = nn.Linear(dims, dims, bias=bias)
        self.key_proj = nn.Linear(dims, dims, bias=bias)
        self.value_proj = nn.Linear(dims, dims, bias=bias)
        self.out_proj = nn.Linear(dims, dims, bias=bias)

    def __call__(self, queries, keys, values, mask=None, cache=None):
        queries = self.query_proj(queries)
        keys = self.key_proj(keys)
        values = self.value_proj(values)

        # Extract some shapes
        num_heads = self.num_heads
        B, L, D = queries.shape

        # Prepare the queries, keys and values for the attention computation
        queries = queries.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)
        keys = keys.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)
        values = values.reshape(B, L, num_heads, -1).transpose(0, 2, 1, 3)

        # Add RoPE to the queries and keys and combine them with the cache
        if cache is not None:
            key_cache, value_cache = cache
            queries = self.rope(queries, offset=key_cache.shape[2])
            keys = self.rope(keys, offset=key_cache.shape[2])
            keys = mx.concatenate([key_cache, keys], axis=2)
            values = mx.concatenate([value_cache, values], axis=2)
        else:
            queries = self.rope(queries)
            keys = self.rope(keys)

        # Finally perform the attention computation
        scale = math.sqrt(1 / queries.shape[-1])
        scores = (queries * scale) @ keys.transpose(0, 1, 3, 2)
        if mask is not None:
            scores = scores + mask
        scores = mx.softmax(scores, axis=-1)
        values_hat = (scores @ values).transpose(0, 2, 1, 3).reshape(B, L, -1)

        # Note that we return the keys and values to possibly be used as a cache
        return self.out_proj(values_hat), (keys, values)


class ParallelBlock(nn.Module):
    def __init__(self, dims: int, num_heads: int, mlp_dims: Optional[int] = None):
        super().__init__()
        mlp_dims = mlp_dims or dims * 4
        self.self_attention = RoPEAttention(dims, num_heads, bias=True)
        self.ln = nn.LayerNorm(dims)
        self.fc1 = nn.Linear(dims, mlp_dims)
        self.fc2 = nn.Linear(mlp_dims, dims)
        self.act = NewGELUActivation()

    def __call__(self, x, x_mask):
        residual = x
        hidden_states = self.ln(x)
        attn_outputs, _ = self.self_attention(
            hidden_states, hidden_states, hidden_states, x_mask
        )
        ff_hidden_states = self.fc2(self.act(self.fc1(hidden_states)))

        hidden_states = attn_outputs + ff_hidden_states + residual

        return hidden_states


class TransformerDecoder(nn.Module):
    def __init__(
        self, num_layers: int, dims: int, num_heads: int, mlp_dims: Optional[int] = None
    ):
        super().__init__()
        self.h = [ParallelBlock(dims, num_heads, mlp_dims) for i in range(num_layers)]

    def __call__(self, x, x_mask):
        for layer in self.h:
            x = layer(x, x_mask)
        return x


class Phi2(nn.Module):
    def __init__(self, config: ModelArgs):
        self.wte = nn.Embedding(config.num_vocab, config.model_dim)
        self.transformer = TransformerDecoder(
            num_layers=config.num_layers,
            dims=config.model_dim,
            num_heads=config.num_heads,
        )

        self.lm_head = LanguageModelingHead(config)

    def __call__(
        self,
        input_ids: mx.array,
        attention_mask: mx.array = None,
    ) -> tuple[mx.array, mx.array]:
        x = self.wte(input_ids)

        if attention_mask is not None:
            # convert 0's to -infs, 1's to 0's, and make it broadcastable
            attention_mask = mx.log(attention_mask)
            attention_mask = mx.expand_dims(attention_mask, (1, 2))
        else:
            attention_mask = nn.MultiHeadAttention.create_additive_causal_mask(
                x.shape[1]
            )

        y = self.transformer(x, attention_mask)
        return self.lm_head(y)

    def generate(self, input_ids, temp=1.0):
        cache = input_ids.tolist()

        # Make an additive causal mask. We will need that to process the prompt.
        mask = nn.MultiHeadAttention.create_additive_causal_mask(input_ids.shape[1])
        mask = mask.astype(self.wte.weight.dtype)

        # First we process the prompt x the same way as in __call__ but
        # save the caches in cache
        x = self.wte(input_ids)
        # for l in self.layers:
        #    x, c = l(x, mask=mask)
        # cache.append(c)  # <--- we store the per layer cache in a
        #      simple python list
        x = self.transformer(x, mask)
        y = self.lm_head(x[:, -1])  # <--- we only care about the last logits
        #      that generate the next token
        y = mx.random.categorical(y * (1 / temp))

        # y now has size [1]
        # Since MLX is lazily evaluated nothing is computed yet.
        # Calling y.item() would force the computation to happen at
        # this point but we can also choose not to do that and let the
        # user choose when to start the computation.
        yield y
        cache += [y.item()]

        # Now we parsed the prompt and generated the first token we
        # need to feed it back into the model and loop to generate the
        # rest.
        while True:
            # Unsqueezing the last dimension to add a sequence length
            # dimension of 1
            x = self.wte(mx.array(cache))
            x = self.transformer(x, mask)
            y = self.lm_head(x[:, -1])
            y = mx.random.categorical(y * (1 / temp))
            cache += [y[0].item()]

            yield y


class LanguageModelingHead(nn.Module):
    def __init__(self, config: ModelArgs) -> None:
        self.ln = nn.LayerNorm(config.model_dim)
        self.linear = nn.Linear(config.model_dim, config.num_vocab)

    def __call__(self, inputs):
        return self.linear(self.ln(inputs))


if __name__ == "__main__":
    model = Phi2(ModelArgs())

    weights = mx.load("weights/phi-2.npz")
    weights = tree_unflatten(list(weights.items()))
    weights = tree_map(lambda p: mx.array(p), weights)

    model.update(weights)

    tokenizer = AutoTokenizer.from_pretrained("microsoft/phi-2", trust_remote_code=True)
    tokens = tokenizer(
        '''def print_prime(n):
   """
   Print all primes between 1 and n
   """''',
        return_tensors="np",
        return_attention_mask=False,
    )

    tokens = {key: mx.array(v) for key, v in tokens.items()}

    print(
        '''def print_prime(n):
   """
   Print all primes between 1 and n
   """'''
    )
    for output in model.generate(**tokens):
        print(tokenizer.decode(output.item()))