"""
Recreate the Core ML model from scratch using
coremltools' neural_network.NeuralNetworkBuilder
"""
import coremltools
import coremltools.models.datatypes as datatypes
from coremltools.models import neural_network as neural_network
from coremltools.models.utils import save_spec
import numpy as np
# get weights
from transformers import GPT2LMHeadModel, GPT2Tokenizer
model_name = "./distilgpt2-base-pretrained-he"
save_directory = "tmp/coreml/"
#!mkdir -p $save_directory
file_name = "model.mlmodel"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
lm_head_model = GPT2LMHeadModel.from_pretrained(model_name).eval()
model = lm_head_model.transformer
wte = model.wte.weight.data.numpy().transpose() # shape (768, 50257) /!\ i hate this
wpe = model.wpe.weight.data.numpy().transpose() # shape (768, 1024)
sequence_length = 64
steps = 6
# build model
input_features = [
('input_ids', datatypes.Array(sequence_length)),
('position_ids', datatypes.Array(sequence_length)),
]
output_features = [('output_logits', None)]
builder = neural_network.NeuralNetworkBuilder(
input_features,
output_features,
mode=None,
disable_rank5_shape_mapping=True,
)
builder.add_expand_dims(
name='input_ids_expanded_to_rank5',
input_name='input_ids',
output_name='input_ids_expanded_to_rank5',
axes=(1, 2, 3, 4)
)
builder.add_expand_dims(
name='position_ids_expanded_to_rank5',
input_name='position_ids',
output_name='position_ids_expanded_to_rank5',
axes=(1, 2, 3, 4)
)
builder.add_embedding(
name='token_embeddings',
input_name='input_ids_expanded_to_rank5',
output_name='token_embeddings',
W=wte,
b=None,
input_dim=50257,
output_channels=768,
has_bias=False,
)
builder.add_embedding(
name='positional_embeddings',
input_name='position_ids_expanded_to_rank5',
output_name='positional_embeddings',
W=wpe,
b=None,
input_dim=1024,
output_channels=768,
has_bias=False,
)
# Input:, Output: (seq, 1, 768, 1, 1)
builder.add_add_broadcastable(
name='embeddings_addition',
input_names=['token_embeddings', 'positional_embeddings'],
output_name=f'{0}_previous_block'
)
for i in range(steps):
print(i)
ln_weight = model.h[i].ln_1.weight.data.numpy().reshape((1, 1, 768, 1, 1))
ln_bias = model.h[i].ln_1.bias.data.numpy().reshape((1, 1, 768, 1, 1))
ln_epsilon = model.h[i].ln_1.eps
builder.add_mvn(
name=f"{i}_block_ln_1",
input_name=f"{i}_previous_block",
# output_name=f"{i}_block_ln_1_output",
output_name=f"{i}_block_ln_1",
across_channels=True,
normalize_variance=True,
epsilon=ln_epsilon
)
builder.add_scale(
name=f"{i}_block_ln_1_scaled",
input_name=f"{i}_block_ln_1",
output_name=f"{i}_block_ln_1_scaled",
W=ln_weight,
b=ln_bias,
has_bias=True,
shape_scale=[768],
shape_bias=[768]
)
builder.add_transpose(
name=f"{i}_block_ln_1_reshape",
input_name=f"{i}_block_ln_1_scaled",
output_name=f"{i}_block_ln_1_scaled_transposed",
axes=(1, 0, 2, 3, 4)
)
conv_1D_bias = model.h[i].attn.c_attn.bias.data.numpy().reshape((1, 1, 2304, 1, 1))
conv_1D_weights = model.h[i].attn.c_attn.weight.data.numpy().transpose().reshape((1, 768, 2304, 1, 1))
builder.add_inner_product(
name=f"{i}_block_attn_conv",
input_name=f"{i}_block_ln_1_scaled_transposed",
output_name=f"{i}_block_attn_conv",
input_channels=768,
output_channels=2304,
W=conv_1D_weights,
b=conv_1D_bias,
has_bias=True
)
builder.add_split(
name=f"{i}_block_attn_qkv_split",
input_name=f"{i}_block_attn_conv",
output_names=[f"{i}_block_attn_q", f"{i}_block_attn_k", f"{i}_block_attn_v"]
)
builder.add_rank_preserving_reshape(
name=f"{i}_block_attn_q_reshape",
input_name=f"{i}_block_attn_q",
output_name=f"{i}_block_attn_q_reshape",
output_shape=(1, 1, sequence_length, 12, 64)
)
builder.add_transpose(
name=f"{i}_block_attn_q_reshape_permuted",
input_name=f"{i}_block_attn_q_reshape",
output_name=f"{i}_block_attn_q_reshape_permuted",
axes=(0, 1, 3, 2, 4)
)
builder.add_rank_preserving_reshape(
name=f"{i}_block_attn_k_reshape",
input_name=f"{i}_block_attn_k",
output_name=f"{i}_block_attn_k_reshape",
output_shape=(1, 1, sequence_length, 12, 64)
)
builder.add_transpose(
name=f"{i}_block_attn_k_reshape_permuted",
input_name=f"{i}_block_attn_k_reshape",
output_name=f"{i}_block_attn_k_reshape_permuted",
axes=(0, 1, 3, 4, 2)
)
builder.add_rank_preserving_reshape(
name=f"{i}_block_attn_v_reshape",
input_name=f"{i}_block_attn_v",
output_name=f"{i}_block_attn_v_reshape",
output_shape=(1, 1, sequence_length, 12, 64)
)
builder.add_transpose(
name=f"{i}_block_attn_v_reshape_permuted",
input_name=f"{i}_block_attn_v_reshape",
output_name=f"{i}_block_attn_v_reshape_permuted",
axes=(0, 1, 3, 2, 4)
)
builder.add_batched_mat_mul(
name=f"{i}_block_attn_qv_matmul",
input_names=[f"{i}_block_attn_q_reshape_permuted", f"{i}_block_attn_k_reshape_permuted"],
output_name=f"{i}_block_attn_qv_matmul"
)
builder.add_scale(
name=f"{i}_block_attn_qv_matmul_scaled",
input_name=f"{i}_block_attn_qv_matmul",
output_name=f"{i}_block_attn_qv_matmul_scaled",
W=np.array(1/8),
b=0,
has_bias=False
)
bias_0 = model.h[i].attn.bias
nd = ns = sequence_length
b = (model.h[i].attn.bias[:, :, ns-nd:ns, :ns]).unsqueeze(0)
builder.add_scale(
name=f"{i}_block_attn_bias",
input_name=f"{i}_block_attn_qv_matmul_scaled",
output_name=f"{i}_block_attn_bias",
W=b,
b=None,
has_bias=False,
shape_scale=[1, sequence_length, sequence_length]
)
bias_constant_0 = - 1e4 * (1 - b)
builder.add_bias(
name=f"{i}_block_attn_afterbias",
input_name=f"{i}_block_attn_bias",
output_name=f"{i}_block_attn_afterbias",
# output_name=f"output_logits",
b=bias_constant_0,
shape_bias=[1, sequence_length, sequence_length],
)
builder.add_squeeze(
name=f"{i}_squeezit",
input_name=f"{i}_block_attn_afterbias",
output_name=f"{i}_squeezit",
axes=[0, 1]
)
builder.add_softmax(
name=f"{i}_block_attn_softmax",
input_name=f"{i}_squeezit",
output_name=f"{i}_block_attn_softmax",
)
builder.add_expand_dims(
name=f"{i}_expandit",
input_name=f"{i}_block_attn_softmax",
output_name=f"{i}_expandit",
axes=[0, 1]
)
builder.add_batched_mat_mul(
name=f"{i}_block_full_attention",
input_names=[f"{i}_expandit", f"{i}_block_attn_v_reshape_permuted"],
output_name=f"{i}_block_full_attention"
)
builder.add_transpose(
name=f"{i}_block_full_attention_merged_t",
input_name=f"{i}_block_full_attention",
output_name=f"{i}_block_full_attention_merged_t",
axes=[0, 1, 3, 2, 4]
)
builder.add_rank_preserving_reshape(
name=f"{i}_block_full_attention_merged",
input_name=f"{i}_block_full_attention_merged_t",
output_name=f"{i}_block_full_attention_merged",
output_shape=[1, 1, 1, sequence_length, 768]
)
builder.add_transpose(
name=f"{i}_block_attn_conv_proj_t",
input_name=f"{i}_block_full_attention_merged",
output_name=f"{i}_block_attn_conv_proj_t",
axes=[0, 3, 4, 1, 2]
)
conv_1D_proj_bias = model.h[i].attn.c_proj.bias.data.numpy().reshape((1, 1, 768, 1, 1))
conv_1D_proj_weights = model.h[i].attn.c_proj.weight.data.numpy().transpose().reshape((1, 768, 768, 1, 1))
# Input:, Output: (1, 3, 768, 1, 1)
builder.add_inner_product(
name=f"{i}_block_attn_conv_proj",
input_name=f"{i}_block_attn_conv_proj_t",
output_name=f"{i}_block_attn_conv_proj",
input_channels=768,
output_channels=768,
W=conv_1D_proj_weights,
b=conv_1D_proj_bias,
has_bias=True
)
# Input: (seq, 1, 768, 1, 1), Output: (1, seq, 768, 1, 1)
builder.add_transpose(
name=f"{i}_previous_block_t",
input_name=f'{i}_previous_block',
output_name=f"{i}_previous_block_t",
axes=[1, 0, 2, 3, 4]
)
# Input: [(1, seq, 768, 1, 1), (1, seq, 768, 1, 1)], Output: (1, seq, 768, 1, 1)
builder.add_add_broadcastable(
name=f"{i}_block_xa_sum",
input_names=[f"{i}_previous_block_t", f"{i}_block_attn_conv_proj"],
output_name=f"{i}_block_xa_sum",
# output_name=f"output_logits"
)
ln_2_weight = model.h[i].ln_2.weight.data.numpy().reshape((1, 1, 768, 1, 1))
ln_2_bias = model.h[i].ln_2.bias.data.numpy().reshape((1, 1, 768, 1, 1))
ln_2_epsilon = model.h[i].ln_2.eps
# Input: (1, seq, 768, 1, 1), Output:
builder.add_mvn(
name=f"{i}_block_ln_2",
input_name=f"{i}_block_xa_sum",
output_name=f"{i}_block_ln_2",
across_channels=True,
normalize_variance=True,
epsilon=ln_2_epsilon
)
builder.add_scale(
name=f"{i}_block_ln_2_scaled",
input_name=f"{i}_block_ln_2",
# output_name=f"output_logits",
output_name=f"{i}_block_ln_2_scaled",
W=ln_2_weight,
b=ln_2_bias,
has_bias=True,
shape_scale=[768],
shape_bias=[768]
)
mlp_conv_1D_fc_bias = model.h[i].mlp.c_fc.bias.data.numpy().reshape((1, 1, 3072, 1, 1))
mlp_conv_1D_fc_weights = model.h[i].mlp.c_fc.weight.data.numpy().transpose().reshape((1, 768, 3072, 1, 1))
# Input:, Output: (1, 3, 3072, 1, 1)
builder.add_inner_product(
name=f"{i}_block_mlp_conv_fc",
input_name=f"{i}_block_ln_2_scaled",
output_name=f"{i}_block_mlp_conv_fc",
# output_name=f"output_logits",
input_channels=768,
output_channels=3072,
W=mlp_conv_1D_fc_weights,
b=mlp_conv_1D_fc_bias,
has_bias=True
)
builder.add_gelu(
name=f"{i}_block_mlp_gelu",
input_name=f"{i}_block_mlp_conv_fc",
output_name=f"{i}_block_mlp_gelu",
# output_name=f"output_logits",
mode='TANH_APPROXIMATION'
)
mlp_conv_1D_proj_bias = model.h[i].mlp.c_proj.bias.data.numpy().reshape((1, 1, 768, 1, 1))
mlp_conv_1D_proj_weights = model.h[i].mlp.c_proj.weight.data.numpy().transpose().reshape((1, 3072, 768, 1, 1))
# Input:, Output: (1, 3, 3072, 1, 1)
builder.add_inner_product(
name=f"{i}_block_mlp_conv_proj",
input_name=f"{i}_block_mlp_gelu",
output_name=f"{i}_block_mlp_conv_proj",
# output_name=f"output_logits",
input_channels=3072,
output_channels=768,
W=mlp_conv_1D_proj_weights,
b=mlp_conv_1D_proj_bias,
has_bias=True
)
builder.add_add_broadcastable(
name=f"{i}_block_xm_sum",
input_names=[f"{i}_block_xa_sum", f"{i}_block_mlp_conv_proj"],
# output_name=f"output_logits"
output_name=f"{i + 1}_previous_block_final"
)
builder.add_transpose(
name=f"{i}_block_xm_sum_t",
input_name=f"{i + 1}_previous_block_final",
output_name=f"{i + 1}_previous_block",
axes=[1, 0, 2, 3, 4]
)
ln_f_weight = model.ln_f.weight.data.numpy().reshape((1, 1, 768, 1, 1))
ln_f_bias = model.ln_f.bias.data.numpy().reshape((1, 1, 768, 1, 1))
ln_f_epsilon = model.ln_f.eps
# Input: (1, seq, 768, 1, 1), Output:
builder.add_mvn(
name=f"ln_f",
input_name=f"{steps}_previous_block_final",
output_name=f"ln_f",
# output_name=f"output_logits",
across_channels=True,
normalize_variance=True,
epsilon=ln_f_epsilon
)
builder.add_scale(
name=f"ln_f_scaled",
input_name=f"ln_f",
output_name=f"ln_f_scaled",
# output_name=f"output_logits",
W=ln_f_weight,
b=ln_f_bias,
has_bias=True,
shape_scale=[768],
shape_bias=[768]
)
lm_head_weights = lm_head_model.lm_head.weight.data.numpy().reshape((1, 50257, 768, 1, 1))
builder.add_inner_product(
name="lm_head",
input_name="ln_f_scaled",
output_name="output_logits",
input_channels=768,
output_channels=50257,
W=lm_head_weights,
b=None,
has_bias=False
)
# compile spec to model
mlmodel = coremltools.models.MLModel(builder.spec)
#save_spec(builder.spec, f'./{model_name}-{sequence_length}-{steps}.mlmodel')
save_spec(builder.spec, f'./{save_directory}{file_name}')
# model = coremltools.models.MLModel('gpt2.mlmodel')
# input_ids = np.zeros(sequence_length)
# position_ids = np.arange(sequence_length).astype(np.float)
# input_data = {
# 'input_ids': input_ids,
# 'position_ids': position_ids,
# }
# predictions = mlmodel.predict(input_data)["output_logits"]
# equal = np.amax(predictions - mlp_conv_proj.detach().numpy())
# print(predictions)
# save_spec(builder.spec, 'gpt2.mlmodel')