Repository: keras-team/keras-nlp

Issue:
from_tensor_slices() method not working with GPT-2 model in gpt2_causal_lm.py example
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.](https://github.com/Neeshamraghav012/keras-nlp/blob/master/keras_nlp/models/gpt2/gpt2_causal_lm.py)
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:
![CI](https://res.cloudinary.com/hire-easy/image/upload/v1678812193/Screenshot_2023-03-14_at_11.24.40_AM_p2ssex.png)

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.


Discussion:
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](https://colab.research.google.com/drive/1-uKhrGp_yvofA-o-5OdStS1IwjnRMQpU?usp=sharing)
Thanks! Would you be interested in making a PR to fix? @TheAthleticCoder 
Yeah Sure! I would like to take this up. Thanks!
@TheAthleticCoder Assigned to you, thanks!
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](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/f_net/f_net_masked_lm.py#L58-L69) and [deberta_v3](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/models/deberta_v3/deberta_v3_masked_lm.py#L64-L75), 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).
cc: @chenmoneygithub Is there any other suggestion or direction I should look at to fix it?
@TheAthleticCoder I think you can just use the previous code snippet to modify the docstring. I tested and it worked well, thanks!

Resolving patch:
diff --git a/keras_nlp/models/gpt2/gpt2_causal_lm.py b/keras_nlp/models/gpt2/gpt2_causal_lm.py
--- a/keras_nlp/models/gpt2/gpt2_causal_lm.py
+++ b/keras_nlp/models/gpt2/gpt2_causal_lm.py
@@ -91,7 +91,7 @@ class GPT2CausalLM(Task):
         "I don't listen to music while coding.",
         "But I watch youtube while coding!",
     ]
-    ds = tf.data.Dataset.from_tensor_slices(features)
+    ds = tf.data.Dataset.from_tensor_slices(features).batch(2)
 
     # Create a `GPT2CausalLM` and fit your data.
     gpt2_lm = keras_nlp.models.GPT2CausalLM.from_preset(
@@ -100,7 +100,7 @@ class GPT2CausalLM(Task):
     gpt2_lm.compile(
         loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
     )
-    gpt2_lm.fit(ds, batch_size=2)
+    gpt2_lm.fit(ds)
     ```
 
     Load a pretrained `GPT2CausalLM` with custom preprocessor, and predict on


Test patch: