Returning all Beams and Probs and adding a Testing Unit #908

TheAthleticCoder · 2023-03-23T14:41:38Z

Attempting to resolve #770 #776

Modifications were made to return all beams and their scores. Also included a test unit.

TheAthleticCoder · 2023-03-23T23:02:09Z

@mattdangerw
Done with the requested changes. I had to change the names of a few variables in the BeamSampler function. Do let me know if I need to rename them to something else 🤔

Also added a unit test function. The tests were quite simple in which I just checked the dimensionality of all the returned prompts and the log_probs. Are there any more tests you would like me to add?
Thanks!

mattdangerw

Thanks! A few comments.

mattdangerw · 2023-03-23T23:08:10Z

keras_nlp/samplers/beam_sampler.py

@@ -77,7 +79,13 @@ def __call__(
        index=0,
        mask=None,
        end_token_id=None,
+        return_all_beams=None,


Do we need this in both places? I think we can probably just take this in at init time.

Is there a workflow that would need this at call time?

Fixed in the latest commit

mattdangerw · 2023-03-23T23:08:27Z

keras_nlp/samplers/beam_sampler.py

    ):
+        if return_all_beams is None:


If we remove the call time argument, we can remove this whole if/else block too

mattdangerw · 2023-03-23T23:12:27Z

keras_nlp/samplers/beam_sampler.py

@@ -65,9 +65,11 @@ def next(prompt, state, index):
    def __init__(


Let's make sure to document this above!

mattdangerw · 2023-03-23T23:23:05Z

keras_nlp/samplers/beam_sampler.py

+        top_beams = tf.math.argmax(all_log_probs, axis=-1)[:, tf.newaxis]
+        prompt = tf.gather(all_prompts, top_beams, axis=1, batch_dims=1)
+
+        if return_all_beams:


I wonder if instead we should do something like this...

By default, return the top beam output sequences with shape (batch_size, length). (as we do currently)

If return_all_beams==True, return an (outputs, log_probs) tuple, with shapes (batch_size, num_beams, length) and (batch_size, num_beams) respectively, where each beam is ordered so the most likely is first.

This would add a little complexity to the implementation (we would probably need to do an argsort and gather for the return_all_beams branch), but it would make the return type more useful and less redundant to the end user.

Want the top beams? That's outputs[:, 0, :]. Want the second to top? That's outputs[:, 1, :].

Sounds good! A concern would be how to handle beams which have the same probability. Ensuring that the output order remains the same no matter how many times we generate it over the input sequence is beneficial.

I think argmax will break ties by going with whatever came first, so I think we could just use it without adding any randomness to the process.

TheAthleticCoder · 2023-03-25T10:44:58Z

@mattdangerw Made the requested changes to the BeamSampler() and to the unit tests. Also edited the documentation. Do let me know if any changes need to be made. Thanks!

chenmoneygithub

Thanks for the PR! Overall looks good to me, left some comments on style.

chenmoneygithub · 2023-03-26T12:12:58Z

keras_nlp/samplers/beam_sampler.py


    Call Args:
        {{call_args}}

    Examples:
+    Example 1:


Instead of Example 1,2,3... let's document what each examples does. For example, you can use "Only return the beam of largest accumulated probability" here, and "Return all beams and their probability" in the next example.

chenmoneygithub · 2023-03-26T12:16:29Z

keras_nlp/samplers/beam_sampler.py

-        top_beams = tf.math.argmax(log_probs, axis=-1)[:, tf.newaxis]
-        prompt = tf.gather(prompt, top_beams, axis=1, batch_dims=1)
-        return tf.squeeze(prompt, axis=1)
+        all_prompts, all_log_probs = unflatten_beams(prompt), unflatten_beams(


Now the line is a bit long, so we can break it down to 2 lines:

all_prompts = unflatten_beams(prompt) all_log_probs = unflatten_beams(log_probs)

chenmoneygithub · 2023-03-26T12:21:47Z

keras_nlp/samplers/beam_sampler.py

+        all_prompts, all_log_probs = unflatten_beams(prompt), unflatten_beams(
+            log_probs
+        )
+        top_beams = tf.math.argmax(all_log_probs, axis=-1)[:, tf.newaxis]


These 2 lines are only useful in the branch if not self.return_all_beams:, we can fold them below to slightly improve our performance.

chenmoneygithub · 2023-03-26T12:27:32Z

keras_nlp/samplers/beam_sampler_test.py

+        self.assertEqual(output[0].shape, (self.batch_size, 5, self.length))
+        self.assertEqual(output[1].shape, (self.batch_size, 5))
+        self.assertTrue(tf.reduce_all(output[1][:, 1:] <= output[1][:, :-1]))
+        self.assertEqual(


Let's also test self.join_as_string(output[0][:, 0, :]) == ["sequentially"] since we are testing returning all beams.

chenmoneygithub · 2023-03-26T12:28:10Z

keras_nlp/samplers/beam_sampler_test.py

+        state_chars = list("sequentially")
+        state = tf.constant([[self.char_lookup[c] for c in state_chars]])
+        prompt = tf.fill((self.batch_size, self.length), self.char_lookup["z"])
+        output = self.sampler_all_beams(


Let's use explicit names: sorted_prompts and sorted_log_probs here so that it's more clear to readers what we are testing.

TheAthleticCoder · 2023-03-26T14:33:14Z

@chenmoneygithub I have made the required changes. Do let me know what else needs to be done. Thanks!

chenmoneygithub

Thanks! Only one minor comment.

chenmoneygithub · 2023-03-28T01:53:35Z

keras_nlp/samplers/beam_sampler.py


    Call Args:
        {{call_args}}

    Examples:
+    1. Return only the beam with the highest accumulated probability.


We can remove the number here.

Hey @chenmoneygithub,
Made the required changes. Do let me know if there are any other changes we should make. Thanks!

TheAthleticCoder added 2 commits March 23, 2023 20:05

Returning all Beams and Probs as well as adding a Testing Unit

9278613

modified test 1

c295fd9

mattdangerw mentioned this pull request Mar 23, 2023

Returns All Beams from Beam Search Utility #776

Closed

TheAthleticCoder added 4 commits March 24, 2023 01:22

overriding changes

ccb7234

added more tests since pipeline works :)

608d558

fixed test cases and variable names

7cc3297

fixed log_prob dimension test

6084f7d

mattdangerw requested changes Mar 23, 2023

View reviewed changes

mattdangerw assigned chenmoneygithub Mar 23, 2023

TheAthleticCoder added 5 commits March 24, 2023 10:36

tried removing call argument

d1c7d90

modified documentation temporarily

726ee8d

sorting the prompts and scores

1248b7f

added batchdim as 1 to fix the extra dim error

6106943

added 1 more test, changed documentation

96f4656

chenmoneygithub suggested changes Mar 26, 2023

View reviewed changes

made the style changes

45e8117

chenmoneygithub approved these changes Mar 28, 2023

View reviewed changes

fixed minor comment

63cf746

TheAthleticCoder requested a review from mattdangerw March 28, 2023 19:49

chenmoneygithub merged commit 58c7e1d into keras-team:master Mar 29, 2023

		@@ -65,9 +65,11 @@ def next(prompt, state, index):
		def __init__(

Returning all Beams and Probs and adding a Testing Unit #908

Returning all Beams and Probs and adding a Testing Unit #908

Uh oh!

Conversation

TheAthleticCoder commented Mar 23, 2023

Uh oh!

TheAthleticCoder commented Mar 23, 2023

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TheAthleticCoder commented Mar 25, 2023

Uh oh!

chenmoneygithub left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

TheAthleticCoder commented Mar 26, 2023

Uh oh!

chenmoneygithub left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!