Port gpt2 transformers checkpoint

[cosmo3769]

Hi @mattdangerw @ariG23498,

Ported GPT2 transformers checkpoint in kerasNLP. Please check. Thank you!

[cosmo3769]

@ariG23498, thank you for your amazing reference repository: ariG23498/keras-nlp-hugging-face-integration 🙏

[mattdangerw]

Thanks! Why do we need to add a new public facing argument hf_key_prefix? Ideally we would keep our exposed API surface minimal.

[ariG23498]

@mattdangerw

I have noticed that some model key names have a prefix that breaks model porting.

If one goes to distill-gpt2 they will find the prefix transformer. The simple prefix will break the model porting code.

With the use of hf_key_prefix I introduce a parameter which can be used to set the prefix while model loading.

An example would be

import keras_nlp

model = keras_nlp.models.GPT2CausalLM.from_preset(
    "hf://distilbert/distilgpt2",
    hf_key_prefix="transformer",
)

print(model.generate(["what is"], max_length=15))

[mattdangerw]

Is there any other prefix we need to check for besides "transformer"? Can we just write some code that either just checks for the "transformer" prefix or any prefix to all the weights? I'm not super familiar with the safetensors API, but there must be a way to list keys right?

I think that'd be a lot more usable, and keep the API clean. In practice I don't think people will actually understand what went wrong, look up the safetensor content, discover the key, and pass it.

[ariG23498]

@mattdangerw I agree completely with your pointers.

Upon talking to Matt we think it would be easiest to apply a regex to capture the prefix if any. The transformers library captures the prefix and then removes it while loading in their models. Unfortunately we would not have that information.

To be precise this is where we would like to apply the regex code
https://github.com/keras-team/keras-nlp/blob/b6877df38d5ddadcd1f7c9c30498b933b4b6ee30/keras_nlp/src/utils/transformers/safetensor_utils.py#L50

I think that would remove the listing keys and checking altogether.

WDYT?

[cosmo3769]

Upon talking to Matt we think it would be easiest to apply a regex to capture the prefix if any.

Yeah, something like this demo. Using regex here to get the prefix upto layer_index: tested with gpt2 model

@ariG23498 @mattdangerw

[mattdangerw]

We should make sure to handle both the sharded and single file safetensor case. I do think we could handle this in SafetensorLoader. Since we either get the full list of keys via file.keys() or safetensor_config["weight_map"].keys() we can use that to resolve the actual key name as needed. Let's try to keep the implementation simple.

[mattdangerw]

Also, please run the formatting script!

[ariG23498]

I like the implementation! WDYT @mattdangerw

My 2 cents:

• We should use better variable naming.
• The name of the function is misleading.
• Adding a docstring here, so that we are well aware of the problem and how to solve it.

[ariG23498]

Come to think of it. I think a better approach to this would be to loop over the keys all at once and have a one-to-one mapping of the hf keys and the keras keys.

Return the map, and then use that map later.

This bypasses the following:

• Looping over all the keys multiple times
• Using this function multiple times

[cosmo3769]

Come to think of it. I think a better approach to this would be to loop over the keys all at once and have a one-to-one mapping of the hf keys and the keras keys.

Makes sense. It will be efficient and real power of this will show up when there will be large number of keys. Mapping it all at once instead of running the loop always surely reduces the time complexity as well.

[mattdangerw]

Yeah caching the prefix in some form sounds good.

Is the prefix always the same for all weights? If so, we could probably do something like this...

def get_prefix(key, dict_like):
    if self.prefix is not None:
        return self.prefix
    keys = dict_like.keys()
    if key in keys:
        self.prefix = ""
    else:
        self.prefix = # Some code to figure out the correct prefix.
    return self.prefix

[mattdangerw]

Looking good! Mostly minor nits, but the comments on the conversion--don't hardcode the query/key/value size, is important.

[mattdangerw]

Thanks! LGTM