Skip to content

Ensure shape in sample so that the shape is correct after TFLite conversion #902

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Mar 29, 2023
Merged

Ensure shape in sample so that the shape is correct after TFLite conversion #902

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e7b691b
trying to fix a shape
chenmoneygithub Mar 21, 2023
d564433
found the issue
chenmoneygithub Mar 22, 2023
8f2c51e
better style
chenmoneygithub Mar 22, 2023
fdef395
trial
chenmoneygithub Mar 28, 2023
af39426
wtf
chenmoneygithub Mar 28, 2023
31e8cff
fix beam sampler
chenmoneygithub Mar 29, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions keras_nlp/samplers/beam_sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,6 +132,9 @@ def body(prompt, state, index, log_probs):
)
beam_indices = indices // vocab_size
next_token = flatten_beams(indices % vocab_size)
# Ensure shape is `[None]`, otherwise it causes issues after
# converting to TFLite.
next_token = tf.ensure_shape(next_token, [None])
# We need `ensure_shape` as `top_k` will change the static shape.
next_log_probs = flatten_beams(next_log_probs)
log_probs = tf.ensure_shape(next_log_probs, log_probs.shape)
Expand Down
4 changes: 3 additions & 1 deletion keras_nlp/samplers/sampler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,11 +110,13 @@ def body(prompt, state, index):
# Compute the softmax distribution for the next token.
logits, state = next(prompt, state, index)
probabilities = keras.activations.softmax(logits)

# Compute the next token.
next_token = self.get_next_token(probabilities)
# Don't overwrite anywhere mask is True.
next_token = tf.cast(next_token, prompt.dtype)
# Ensure shape is `[None]`, otherwise it causes issues after
# converting to TFLite.
next_token = tf.ensure_shape(next_token, [None])
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you have an example of where this is causing a failure? Very curious why this is needed.

Is there someway we can test the failure without needing to depend on TFLite you think?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very good question - I am going to talk to TFLite folks after I get back, I know there is some effort on bringing TFLite testing to TF/Keras before conversion.

For the issue I saw - without this ensure_shape, next_token becomes shape [] instead of [None], which causes dynamic_update_slice to have a shape mismatch between operand (1st arg) and update (2nd arg).

next_token = tf.where(mask[:, index], prompt[:, index], next_token)
# Update the prompt with the next token.
next_token = next_token[:, tf.newaxis]
Expand Down