Skip to content

Commit

Permalink
indentation fix
Browse files Browse the repository at this point in the history
  • Loading branch information
sineeli committed Jun 7, 2024
1 parent 256ef8d commit b30dbba
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 60 deletions.
1 change: 0 additions & 1 deletion keras_nlp/src/models/llama/llama_presets.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,5 +45,4 @@
},
"kaggle_handle": "kaggle://sineeli/vicuna/keras/vicuna_1.5_7b_en/1",
},

}
102 changes: 43 additions & 59 deletions tools/checkpoint_conversion/convert_llama_checkpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,20 +18,20 @@
import tempfile
import traceback

os.environ['KERAS_BACKEND'] = "torch"
os.environ["KERAS_BACKEND"] = "torch"

import numpy as np
import torch
from absl import app
from absl import flags
from keras import ops
from transformers import AutoTokenizer
from transformers import LlamaForCausalLM
import numpy as np # noqa: E402
import torch # noqa: E402
from absl import app # noqa: E402
from absl import flags # noqa: E402
from keras import ops # noqa: E402
from transformers import AutoTokenizer # noqa: E402
from transformers import LlamaForCausalLM # noqa: E402

from keras_nlp import upload_preset
from keras_nlp.models import LlamaBackbone
from keras_nlp.models import LlamaCausalLMPreprocessor
from keras_nlp.models import LlamaTokenizer
from keras_nlp import upload_preset # noqa: E402
from keras_nlp.models import LlamaBackbone # noqa: E402
from keras_nlp.models import LlamaCausalLMPreprocessor # noqa: E402
from keras_nlp.models import LlamaTokenizer # noqa: E402

PRESET_MAP = {
"llama2_7b_en": "meta-llama/Llama-2-7b-hf",
Expand All @@ -51,18 +51,20 @@
)

flags.DEFINE_string(
name="validate_dtype",
default='bfloat16',
help=("The dtype of the two models while validating numerics."
"can be 'float32', 'float16', or 'bfloat16'"
name="validate_dtype",
default="bfloat16",
help=(
"The dtype of the two models while validating numerics."
"can be 'float32', 'float16', or 'bfloat16'"
),
)

flags.DEFINE_string(
name="save_dtype",
default='bfloat16',
help=("The dtype of the two models while validating numerics."
"can be 'float32', 'float16', or 'bfloat16'"
name="save_dtype",
default="bfloat16",
help=(
"The dtype of the two models while validating numerics."
"can be 'float32', 'float16', or 'bfloat16'"
),
)

Expand All @@ -74,9 +76,10 @@
"The link to upload the model. can be in these formats: "
"`kaggle://<KAGGLE_USERNAME>/<MODEL>/<FRAMEWORK>/<VARIATION>`, "
"`hf://[<HF_USERNAME>/]<MODEL>`"
)
),
)


def convert_checkpoints(keras_nlp_model, hf_model):
config = hf_model.config

Expand All @@ -89,8 +92,7 @@ def convert_checkpoints(keras_nlp_model, hf_model):
i
]._self_attention_layer._key_dense.set_weights(
[
hf_model.model.layers[i]
.self_attn.k_proj.weight.T.reshape(
hf_model.model.layers[i].self_attn.k_proj.weight.T.reshape(
config.hidden_size,
config.num_key_value_heads,
config.hidden_size // config.num_attention_heads,
Expand All @@ -101,8 +103,7 @@ def convert_checkpoints(keras_nlp_model, hf_model):
i
]._self_attention_layer._query_dense.set_weights(
[
hf_model.model.layers[i]
.self_attn.q_proj.weight.T.reshape(
hf_model.model.layers[i].self_attn.q_proj.weight.T.reshape(
config.hidden_size,
config.num_attention_heads,
config.hidden_size // config.num_attention_heads,
Expand All @@ -113,8 +114,7 @@ def convert_checkpoints(keras_nlp_model, hf_model):
i
]._self_attention_layer._value_dense.set_weights(
[
hf_model.model.layers[i]
.self_attn.v_proj.weight.T.reshape(
hf_model.model.layers[i].self_attn.v_proj.weight.T.reshape(
config.hidden_size,
config.num_key_value_heads,
config.hidden_size // config.num_attention_heads,
Expand All @@ -125,8 +125,7 @@ def convert_checkpoints(keras_nlp_model, hf_model):
i
]._self_attention_layer._output_dense.set_weights(
[
hf_model.model.layers[i]
.self_attn.o_proj.weight.T.reshape(
hf_model.model.layers[i].self_attn.o_proj.weight.T.reshape(
config.num_attention_heads,
config.hidden_size // config.num_attention_heads,
config.hidden_size,
Expand All @@ -136,47 +135,30 @@ def convert_checkpoints(keras_nlp_model, hf_model):
keras_nlp_model.transformer_layers[
i
]._self_attention_layernorm.set_weights(
[
hf_model.model.layers[i]
.input_layernorm.weight
]
[hf_model.model.layers[i].input_layernorm.weight]
)
keras_nlp_model.transformer_layers[
i
]._feedforward_intermediate_dense.set_weights(
[
hf_model.model.layers[i]
.mlp.up_proj.weight.T
]
[hf_model.model.layers[i].mlp.up_proj.weight.T]
)
keras_nlp_model.transformer_layers[
i
]._feedforward_output_dense.set_weights(
[
hf_model.model.layers[i]
.mlp.down_proj.weight.T
]
[hf_model.model.layers[i].mlp.down_proj.weight.T]
)
keras_nlp_model.transformer_layers[
i
]._feedforward_gate_dense.set_weights(
[
hf_model.model.layers[i]
.mlp.gate_proj.weight.T
]
[hf_model.model.layers[i].mlp.gate_proj.weight.T]
)
keras_nlp_model.transformer_layers[
i
]._feedforward_layernorm.set_weights(
[
hf_model.model.layers[i]
.post_attention_layernorm.weight.detach()
]
[hf_model.model.layers[i].post_attention_layernorm.weight.detach()]
)

keras_nlp_model.layer_norm.set_weights(
[hf_model.model.norm.weight]
)
keras_nlp_model.layer_norm.set_weights([hf_model.model.norm.weight])
keras_nlp_model.token_embedding.reverse_embeddings.assign(
hf_model.lm_head.weight.T
)
Expand Down Expand Up @@ -240,18 +222,20 @@ def main(_):
upload_link = FLAGS.upload_link
hf_preset = PRESET_MAP[preset]
torch_dtype = torch_dtype_map.get(FLAGS.validate_dtype)

# === Create the temporary save directories ===
temp_dir = tempfile.mkdtemp()

try:
# === Load the Huggingface model ===
hf_model = LlamaForCausalLM.from_pretrained(
hf_preset, torch_dtype=torch_dtype
)
hf_tokenizer = AutoTokenizer.from_pretrained(hf_preset)
hf_model.eval()
print(f"\n-> Huggingface model and tokenizer loaded with dtype: {FLAGS.validate_dtype}")
print(
f"\n-> Huggingface model and tokenizer loaded with dtype: {FLAGS.validate_dtype}"
)

# === Load the KerasNLP model ===
backbone_kwargs = dict(
Expand Down Expand Up @@ -280,18 +264,18 @@ def main(_):
test_tokenizer(keras_nlp_tokenizer, hf_tokenizer)
test_model(keras_nlp_model, keras_nlp_tokenizer, hf_model, hf_tokenizer)
print("\n-> Tests passed!")

keras_nlp_model.save_weights(os.path.join(temp_dir, "model.weights.h5"))
print(f"\n-> Saved the model weights in {FLAGS.validate_dtype}")

del keras_nlp_model, hf_model
gc.collect()

# === Save the weights again in user defined dtype ===
backbone_kwargs["dtype"] = FLAGS.save_dtype
keras_nlp_model = LlamaBackbone(**backbone_kwargs)
keras_nlp_model.load_weights(os.path.join(temp_dir, "model.weights.h5"))

# === Save the model ===
keras_nlp_model.save_to_preset(preset)
print(f"\n-> Saved the model preset in {FLAGS.save_dtype}")
Expand All @@ -310,4 +294,4 @@ def main(_):

if __name__ == "__main__":
flags.mark_flag_as_required("preset")
app.run(main)
app.run(main)

0 comments on commit b30dbba

Please sign in to comment.