from_tensor_slices() method not working with GPT-2 model in gpt2_causal_lm.py example

[Neeshamraghav012]

Hi team,

I was doing some experiments with gpt2 causal lm, I started by running the examples that are given in the documentation
Link to documentation.
But this example didn't worked,

features = [
        "I don't listen to music while coding.",
        "But I watch youtube while coding!",
]

ds = tf.data.Dataset.from_tensor_slices(features)

gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
        "gpt2_base_en",
)
gpt2_lm.compile(
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
gpt2_lm.fit(ds, batch_size=2)

Output:

But when I passed the features directly it worked. Here is the working code;

features = [
        "I don't listen to music while coding.",
        "But I watch youtube while coding!",
]

ds = tf.data.Dataset.from_tensor_slices(features)

gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
        "gpt2_base_en",
)
gpt2_lm.compile(
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
gpt2_lm.fit(features, batch_size=2)

It seems that there might be an issue with the from_tensor_slices() method not working properly with the GPT-2 model in this example. I wanted to report this issue and bring it to the attention of the KerasNLP community. I am very much interested in solving these kinds of bugs that are present in existing pre-trained models.

[TheAthleticCoder]

I would like to add that instead of passing features in gpt2_lm.fit(features, batch_size=2) we can pass ds but with a slight modification:

features = [
        "I don't listen to music while coding.",
        "But I watch youtube while coding!",
]

ds = tf.data.Dataset.from_tensor_slices(features)

gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
        "gpt2_base_en",
)
gpt2_lm.compile(
        loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
)
new_ds = ds.batch(2)
gpt2_lm.fit(new_ds)

It executes successfully as can be seen in this colab link

[chenmoneygithub]

Thanks! Would you be interested in making a PR to fix? @TheAthleticCoder

[TheAthleticCoder]

Yeah Sure! I would like to take this up. Thanks!

[chenmoneygithub]

@TheAthleticCoder Assigned to you, thanks!

[TheAthleticCoder]

So, we aren't allowed to pass <TensorSliceDataset element_spec=TensorSpec(shape=(), dtype=tf.string, name=None)> to model.fit() since it is not directly compatible. So there are two ways to approach this:

1. Similar to other LM documentation such as: f_net and deberta_v3, we can just mention in gpt2's documentation to simply pass features into gpt2_lm.fit() and not use .from_tensor_slices(features)

2. Or, we can say that we batch ds outside using ds.batch(batch_size) and then pass it like gpt2_lm.fit(ds) (as I experimented in the previous comment).

[TheAthleticCoder]

cc: @chenmoneygithub Is there any other suggestion or direction I should look at to fix it?

[chenmoneygithub]

@TheAthleticCoder I think you can just use the previous code snippet to modify the docstring. I tested and it worked well, thanks!