Adding a temperature argument to the base sampler class and related tests #951
Conversation
keras_nlp/samplers/top_k_sampler.py
Outdated
| self, | ||
| k=5, | ||
| seed=None, | ||
| temperature=1.0, |
There was a problem hiding this comment.
if you do the changes suggested above, you won't have to do anything beyond passing a **kwargs through.
|
|
||
| self.next = next | ||
| self.sampler = TopKSampler(k=5) | ||
| self.sampler_temperature = TopKSampler(k=5, temperature=1e-5) |
There was a problem hiding this comment.
Just define this in the test below, as it is only used for one test.
| next=next, | ||
| prompt=prompt, | ||
| ) | ||
| output_ids = set(output[0].numpy()) |
There was a problem hiding this comment.
no need to do set here, just self.assertAllEqual(output, [0] * self.length) or something like that.
|
/gcbrun |
There was a problem hiding this comment.
This looks great! I think we just need to merge this with the latest changes on master.
For beam search, top-k, greedy, contrastive, we can probably just very pass temperature=1.0 in setup for the test sampler, as the output won't change for these "top" samplers. That way our test files stay short.
For top-p and random we can include a simple test that assert that a very low temp will lead to deterministic output.
| output_ids = set(output[0].numpy()) | ||
| self.assertContainsSubset(output_ids, range(3)) | ||
|
|
||
| def test_outputs_in_top_p_with_temperature(self): |
There was a problem hiding this comment.
I think we could make this simpler, we are just trying to test that temperature is working. Something like (haven't run this)...
def test_temperature(self):
def next(prompt, state, index):
# Return a distribution where each id is progressively less likely.
logits = tf.range(self.vocab_size, 0, -1, dtype="float32")
logits = tf.repeat(logits[tf.newaxis, :], self.batch_size, axis=0)
return tf.constant(logits), state
prompt = tf.fill((self.batch_size, self.length), self.char_lookup["z"])
# Use a very low temperature, so we always sample the first token.
output = TopPSampler(p=0.5, temperature=1e-6)(
next=next,
prompt=prompt,
)
self.assertAllEqual(output, tf.zero_like(prompt))There was a problem hiding this comment.
This works for the top_p_sampler. However, when I apply the same for the random_sampler, it only passes 75% of the tests as seen here.
Not sure why.
|
Thanks! Will resolve merge conflicts as I merge this. |
|
@mattdangerw Should I add the temperature arg to the contrastive sampler as well in this thread? Or will we be opening another issue for the same? |
Follow up is fine, I will merge this now. |
|
/gcbrun |
Partially resolves #947
Currently working on
sampler.pyand added a small unit test totop_kfile.