top p search and testing #233
Conversation
keras_nlp/utils/text_generation.py
Outdated
| pred = tf.keras.activations.softmax(pred, axis=-1) | ||
| # Sort preds in descending order. | ||
| sorted_preds, sorted_indices = tf.math.top_k( | ||
| pred, k=pred.shape[-1], sorted=True |
There was a problem hiding this comment.
Let's be consistent with -1 or 1 here.
keras_nlp/utils/text_generation.py
Outdated
| replaced with `pad_token_id`. | ||
| pad_token_id: int, defaults to 0. The pad token after `end_token_id` | ||
| is received. | ||
| filter_value: float, defaults to -Inf. The value for filtering out |
There was a problem hiding this comment.
Do we want to allow customization on filter_value? The code uses filter_value to set certain tokens to have probability 0?
There was a problem hiding this comment.
I don't think we need to allow customization, we can just mask out tokens that are unused by having filter_value = 0
There was a problem hiding this comment.
Yea, in that case, let's remove this argument.
keras_nlp/utils/text_generation.py
Outdated
| ) | ||
| # Filter out unmasked tokens and sample from filtered distribution. | ||
| probs = tf.where( | ||
| shifted_keep_mask, sorted_preds, tf.fill(pred.shape, filter_value) |
There was a problem hiding this comment.
Should filter_value always be 0? I saw in the docstring it defaults to -Inf.
There was a problem hiding this comment.
Yep filter_value should be 0, edited
| self.assertAllEqual(output_logit, output_probs) | ||
|
|
||
|
|
||
| class TopPSamplingTextGenerationTest(tf.test.TestCase): |
There was a problem hiding this comment.
Let's keep test name and method name the same - TopPSearchTextGenerationTest
There was a problem hiding this comment.
edited for this and other testing names
keras_nlp/utils/text_generation.py
Outdated
| replaced with `pad_token_id`. | ||
| pad_token_id: int, defaults to 0. The pad token after `end_token_id` | ||
| is received. | ||
| filter_value: float, defaults to -Inf. The value for filtering out |
There was a problem hiding this comment.
Yea, in that case, let's remove this argument.
| prob = tf.constant([[0.4, 0.3, 0.2, 0.1]]) | ||
| return tf.repeat(prob, batch_size, axis=0) | ||
|
|
||
| # Test that it only samples from tokens that sum up to p |
There was a problem hiding this comment.
nit: period at end.
keras_nlp/utils/text_generation.py
Outdated
| ) | ||
| if p <= 0 or p >= 1: | ||
| raise ValueError( | ||
| "p should be a float strictly between 0 and 1 (0 < p < 1)." |
There was a problem hiding this comment.
surround arg names in backticks (p) also show what value we received, something like.
`p` should be in the range (0, 1). Received: `p={p}`.
keras_nlp/utils/text_generation.py
Outdated
| """ | ||
| Text generation utility based on top-p (nucleus) sampling. | ||
| Top-p search filters the top token with probabilities that sum up to p, and |
There was a problem hiding this comment.
maybe we can explain this a little more clearly, also remember to backtick argument names so it's clear what they are...
Top-p search selects tokens from the smallest subset of output probabilities that sum to greater than p. Put another way, top-p will first order token predictions by likelihood, and ignore all tokens after the cumulative probably of selected tokens exceeds p.
| prompt = tf.fill((BATCH_SIZE, 1), START_ID) | ||
| # Print the generated sequence (token ids). |
No description provided.