We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Return All Beams from Beam Search Utility
Follow-up to https://github.com/keras-team/keras-nlp/issues/750#issuecomment-1439963933

- Return all beams (tensor of shape `(batch_size, num_beams, sequence_length)`).
- Return the beam scores as well (tensor of shape `(batch_size, num_beams)`).

The changes are to be made to this utility function: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/utils/text_generation.py#L213.

@mattdangerw, do we want to sort the outputs based on log prob and return? Or is that unnecessary overhead?

Some basic tinkering done because a user wanted a quick fix: https://github.com/abheesht17/keras-nlp/blob/beam_search_return_all_seqs/keras_nlp/utils/text_generation.py#L400-L414

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `keras_nlp/samplers/beam_sampler.py`
Content:
```
1 # Copyright 2023 The KerasNLP Authors
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 """Beam Sampler."""
15 
16 import tensorflow as tf
17 from tensorflow import keras
18 from tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice
19 
20 from keras_nlp.api_export import keras_nlp_export
21 from keras_nlp.samplers.sampler import Sampler
22 from keras_nlp.samplers.sampler import call_args_docstring
23 from keras_nlp.utils.python_utils import format_docstring
24 
25 
26 @format_docstring(call_args=call_args_docstring)
27 @keras_nlp_export("keras_nlp.samplers.BeamSampler")
28 class BeamSampler(Sampler):
29     """Beam Sampler class.
30 
31     This sampler implements beam search algorithm. At each time-step, beam
32     search keeps the beams (sequences) of the top `num_beams` highest
33     accumulated probabilities, and uses each one of the beams to predict
34     candidate next tokens.
35 
36     Args:
37         num_beams: int. The number of beams that should be kept at each
38             time-step. `num_beams` should be strictly positive.
39 
40     Call Args:
41         {{call_args}}
42 
43     Examples:
44     ```python
45     # Use a simple alphabet of lowercase characters to [0, 26).
46     int_lookup = {i: chr(i + ord('a')) for i in range(26)}
47     char_lookup = {v: k for k, v in int_lookup.items()}
48     batch_size, length, vocab_size = 1, 12, len(int_lookup)
49 
50     def next(prompt, state, index):
51         # A uniform distribution over our alphabet.
52         logits = tf.ones((batch_size, vocab_size))
53         return logits, state
54 
55     output = keras_nlp.samplers.BeamSampler()(
56         next=next,
57         prompt=tf.fill((batch_size, length,), char_lookup['z']),
58         index=5,
59     )
60     print(["".join([int_lookup[i] for i in s]) for s in output.numpy()])
61     # >>> "zzzzzaaaaaaa"
62     ```
63     """
64 
65     def __init__(
66         self,
67         num_beams=5,
68     ):
69         super().__init__()
70         self.num_beams = num_beams
71 
72     def __call__(
73         self,
74         next,
75         prompt,
76         state=None,
77         index=0,
78         mask=None,
79         end_token_id=None,
80     ):
81         batch_size, max_length = tf.shape(prompt)[0], tf.shape(prompt)[1]
82         # Make sure max length and start index are the same dtype.
83         index = tf.cast(index, max_length.dtype)
84 
85         def create_beams(x):
86             """Add initial beam state."""
87             return tf.repeat(x, self.num_beams, axis=0)
88 
89         def flatten_beams(x):
90             """Combine the beam dim and batch dim."""
91             flat_shape = [batch_size * self.num_beams] + x.shape.as_list()[2:]
92             return tf.reshape(x, shape=flat_shape)
93 
94         def unflatten_beams(x):
95             """Separate the beam dim and batch dim."""
96             unflat_shape = [batch_size, self.num_beams] + x.shape.as_list()[1:]
97             return tf.reshape(x, shape=unflat_shape)
98 
99         mask = tf.zeros_like(prompt, dtype=tf.bool) if mask is None else mask
100         # `tf.while_loop` will not accept `None` as a value for `loop_vars`.
101         state = () if state is None else state
102         # Add extra sequences for each beam.
103         prompt, mask = create_beams(prompt), create_beams(mask)
104         state = tf.nest.map_structure(create_beams, state)
105         # Setup the initial beam log-likelihoods.
106         # On the first loop, make sure only the original beam is considered.
107         log_probs = tf.constant([[0.0] + [-1e9] * (self.num_beams - 1)])
108         log_probs = flatten_beams(tf.repeat(log_probs, batch_size, axis=0))
109 
110         def cond(prompt, state, index, log_probs):
111             if end_token_id is None:
112                 return True
113             # Stop if all sequences have produced a *new* end_token_id.
114             end_tokens = (prompt == end_token_id) & (~mask)
115             prompt_done = tf.reduce_any(end_tokens, axis=-1)
116             return not tf.reduce_all(prompt_done)
117 
118         def body(prompt, state, index, log_probs):
119             # Compute the softmax distribution for the next token.
120             logits, state = next(prompt, state, index)
121             vocab_size = tf.shape(logits)[-1]
122             probs = keras.activations.softmax(logits)
123 
124             # Compute the running log-likelihood of each new candidate.
125             next_log_probs = tf.math.log(probs) + log_probs[..., tf.newaxis]
126             # Reshape `preds` to shape `(batch_size, num_beams * vocab_size)`.
127             next_log_probs = tf.reshape(next_log_probs, shape=[batch_size, -1])
128 
129             # Compute the top beam indices and next tokens.
130             next_log_probs, indices = tf.math.top_k(
131                 next_log_probs, k=self.num_beams, sorted=False
132             )
133             beam_indices = indices // vocab_size
134             next_token = flatten_beams(indices % vocab_size)
135             # We need `ensure_shape` as `top_k` will change the static shape.
136             next_log_probs = flatten_beams(next_log_probs)
137             log_probs = tf.ensure_shape(next_log_probs, log_probs.shape)
138 
139             def gather_beams(x):
140                 x = unflatten_beams(x)
141                 x = tf.gather(x, beam_indices, axis=1, batch_dims=1)
142                 return flatten_beams(x)
143 
144             prompt = gather_beams(prompt)
145             state = tf.nest.map_structure(gather_beams, state)
146 
147             # Update each beam with the next token.
148             next_token = tf.cast(next_token, prompt.dtype)
149             # Don't overwrite anywhere mask is True.
150             next_token = tf.where(mask[:, index], prompt[:, index], next_token)
151             # Update the prompt with the next token.
152             next_token = next_token[:, tf.newaxis]
153             prompt = dynamic_update_slice(prompt, next_token, [0, index])
154             # Return the iteration of the loop state.
155             return (prompt, state, index + 1, log_probs)
156 
157         prompt, _, _, log_probs = tf.while_loop(
158             cond=cond,
159             body=body,
160             loop_vars=(prompt, state, index, log_probs),
161             maximum_iterations=(max_length - index),
162         )
163 
164         # Gather the top beam at each batch index.
165         prompt, log_probs = unflatten_beams(prompt), unflatten_beams(log_probs)
166         top_beams = tf.math.argmax(log_probs, axis=-1)[:, tf.newaxis]
167         prompt = tf.gather(prompt, top_beams, axis=1, batch_dims=1)
168         return tf.squeeze(prompt, axis=1)
169 
170     def get_config(self):
171         config = super().get_config()
172         config.update(
173             {
174                 "num_beams": self.num_beams,
175             }
176         )
177         return config
178
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/keras_nlp/samplers/beam_sampler.py b/keras_nlp/samplers/beam_sampler.py
--- a/keras_nlp/samplers/beam_sampler.py
+++ b/keras_nlp/samplers/beam_sampler.py
@@ -36,11 +36,14 @@
     Args:
         num_beams: int. The number of beams that should be kept at each
             time-step. `num_beams` should be strictly positive.
+        return_all_beams: bool. When set to `True`, the sampler will return the top prompt,
+            all prompts and their respective probabilities score.
 
     Call Args:
         {{call_args}}
 
     Examples:
+    Return only the beam with the highest accumulated probability.
     ```python
     # Use a simple alphabet of lowercase characters to [0, 26).
     int_lookup = {i: chr(i + ord('a')) for i in range(26)}
@@ -60,14 +63,41 @@
     print(["".join([int_lookup[i] for i in s]) for s in output.numpy()])
     # >>> "zzzzzaaaaaaa"
     ```
+    Return all beams and their probabilities.
+    ```python
+    # Use a simple alphabet of lowercase characters to [0, 26).
+    int_lookup = {i: chr(i + ord('a')) for i in range(26)}
+    char_lookup = {v: k for k, v in int_lookup.items()}
+    batch_size, length, vocab_size = 1, 12, len(int_lookup)
+
+    def next(prompt, state, index):
+        # A uniform distribution over our alphabet.
+        logits = tf.ones((batch_size, vocab_size))
+        return logits, state
+
+    output = keras_nlp.samplers.BeamSampler(return_all_beams=True)(
+        next=next,
+        prompt=tf.fill((batch_size, length,), char_lookup['z']),
+        index=5,
+    )
+
+    print(output[0].shape)
+    # >>> (1, 5, 12)
+    print(output[1].shape)
+    # >>> (1, 5)
+    print(["".join([int_lookup[i] for i in s]) for s in output[0][0].numpy()])
+    # >>> "zzzzzaaaaaaa"
+    ```
     """
 
     def __init__(
         self,
         num_beams=5,
+        return_all_beams=False,
     ):
         super().__init__()
         self.num_beams = num_beams
+        self.return_all_beams = return_all_beams
 
     def __call__(
         self,
@@ -161,17 +191,32 @@
             maximum_iterations=(max_length - index),
         )
 
-        # Gather the top beam at each batch index.
-        prompt, log_probs = unflatten_beams(prompt), unflatten_beams(log_probs)
-        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 = unflatten_beams(prompt)
+        all_log_probs = unflatten_beams(log_probs)
+
+        if self.return_all_beams:
+            sorted_indices = tf.argsort(
+                all_log_probs, axis=-1, direction="DESCENDING"
+            )
+            sorted_log_probs = tf.gather(
+                all_log_probs, sorted_indices, axis=-1, batch_dims=1
+            )
+            sorted_prompts = tf.gather(
+                all_prompts, sorted_indices, axis=1, batch_dims=1
+            )
+            return sorted_prompts, sorted_log_probs
+        else:
+            # Gather the top beam at each batch index.
+            top_beams = tf.math.argmax(all_log_probs, axis=-1)[:, tf.newaxis]
+            prompt = tf.gather(all_prompts, top_beams, axis=1, batch_dims=1)
+            return tf.squeeze(prompt, axis=1)
 
     def get_config(self):
         config = super().get_config()
         config.update(
             {
                 "num_beams": self.num_beams,
+                "return_all_beams": self.return_all_beams,
             }
         )
         return config

{"golden_diff": "diff --git a/keras_nlp/samplers/beam_sampler.py b/keras_nlp/samplers/beam_sampler.py\n--- a/keras_nlp/samplers/beam_sampler.py\n+++ b/keras_nlp/samplers/beam_sampler.py\n@@ -36,11 +36,14 @@\n     Args:\n         num_beams: int. The number of beams that should be kept at each\n             time-step. `num_beams` should be strictly positive.\n+        return_all_beams: bool. When set to `True`, the sampler will return the top prompt,\n+            all prompts and their respective probabilities score.\n \n     Call Args:\n         {{call_args}}\n \n     Examples:\n+    Return only the beam with the highest accumulated probability.\n     ```python\n     # Use a simple alphabet of lowercase characters to [0, 26).\n     int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n@@ -60,14 +63,41 @@\n     print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n     # >>> \"zzzzzaaaaaaa\"\n     ```\n+    Return all beams and their probabilities.\n+    ```python\n+    # Use a simple alphabet of lowercase characters to [0, 26).\n+    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n+    char_lookup = {v: k for k, v in int_lookup.items()}\n+    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n+\n+    def next(prompt, state, index):\n+        # A uniform distribution over our alphabet.\n+        logits = tf.ones((batch_size, vocab_size))\n+        return logits, state\n+\n+    output = keras_nlp.samplers.BeamSampler(return_all_beams=True)(\n+        next=next,\n+        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n+        index=5,\n+    )\n+\n+    print(output[0].shape)\n+    # >>> (1, 5, 12)\n+    print(output[1].shape)\n+    # >>> (1, 5)\n+    print([\"\".join([int_lookup[i] for i in s]) for s in output[0][0].numpy()])\n+    # >>> \"zzzzzaaaaaaa\"\n+    ```\n     \"\"\"\n \n     def __init__(\n         self,\n         num_beams=5,\n+        return_all_beams=False,\n     ):\n         super().__init__()\n         self.num_beams = num_beams\n+        self.return_all_beams = return_all_beams\n \n     def __call__(\n         self,\n@@ -161,17 +191,32 @@\n             maximum_iterations=(max_length - index),\n         )\n \n-        # Gather the top beam at each batch index.\n-        prompt, log_probs = unflatten_beams(prompt), unflatten_beams(log_probs)\n-        top_beams = tf.math.argmax(log_probs, axis=-1)[:, tf.newaxis]\n-        prompt = tf.gather(prompt, top_beams, axis=1, batch_dims=1)\n-        return tf.squeeze(prompt, axis=1)\n+        all_prompts = unflatten_beams(prompt)\n+        all_log_probs = unflatten_beams(log_probs)\n+\n+        if self.return_all_beams:\n+            sorted_indices = tf.argsort(\n+                all_log_probs, axis=-1, direction=\"DESCENDING\"\n+            )\n+            sorted_log_probs = tf.gather(\n+                all_log_probs, sorted_indices, axis=-1, batch_dims=1\n+            )\n+            sorted_prompts = tf.gather(\n+                all_prompts, sorted_indices, axis=1, batch_dims=1\n+            )\n+            return sorted_prompts, sorted_log_probs\n+        else:\n+            # Gather the top beam at each batch index.\n+            top_beams = tf.math.argmax(all_log_probs, axis=-1)[:, tf.newaxis]\n+            prompt = tf.gather(all_prompts, top_beams, axis=1, batch_dims=1)\n+            return tf.squeeze(prompt, axis=1)\n \n     def get_config(self):\n         config = super().get_config()\n         config.update(\n             {\n                 \"num_beams\": self.num_beams,\n+                \"return_all_beams\": self.return_all_beams,\n             }\n         )\n         return config\n", "issue": "Return All Beams from Beam Search Utility\nFollow-up to https://github.com/keras-team/keras-nlp/issues/750#issuecomment-1439963933\r\n\r\n- Return all beams (tensor of shape `(batch_size, num_beams, sequence_length)`).\r\n- Return the beam scores as well (tensor of shape `(batch_size, num_beams)`).\r\n\r\nThe changes are to be made to this utility function: https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/utils/text_generation.py#L213.\r\n\r\n@mattdangerw, do we want to sort the outputs based on log prob and return? Or is that unnecessary overhead?\r\n\r\nSome basic tinkering done because a user wanted a quick fix: https://github.com/abheesht17/keras-nlp/blob/beam_search_return_all_seqs/keras_nlp/utils/text_generation.py#L400-L414\n", "before_files": [{"content": "# Copyright 2023 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Beam Sampler.\"\"\"\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice\n\nfrom keras_nlp.api_export import keras_nlp_export\nfrom keras_nlp.samplers.sampler import Sampler\nfrom keras_nlp.samplers.sampler import call_args_docstring\nfrom keras_nlp.utils.python_utils import format_docstring\n\n\n@format_docstring(call_args=call_args_docstring)\n@keras_nlp_export(\"keras_nlp.samplers.BeamSampler\")\nclass BeamSampler(Sampler):\n    \"\"\"Beam Sampler class.\n\n    This sampler implements beam search algorithm. At each time-step, beam\n    search keeps the beams (sequences) of the top `num_beams` highest\n    accumulated probabilities, and uses each one of the beams to predict\n    candidate next tokens.\n\n    Args:\n        num_beams: int. The number of beams that should be kept at each\n            time-step. `num_beams` should be strictly positive.\n\n    Call Args:\n        {{call_args}}\n\n    Examples:\n    ```python\n    # Use a simple alphabet of lowercase characters to [0, 26).\n    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n    char_lookup = {v: k for k, v in int_lookup.items()}\n    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n\n    def next(prompt, state, index):\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((batch_size, vocab_size))\n        return logits, state\n\n    output = keras_nlp.samplers.BeamSampler()(\n        next=next,\n        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n        index=5,\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> \"zzzzzaaaaaaa\"\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        num_beams=5,\n    ):\n        super().__init__()\n        self.num_beams = num_beams\n\n    def __call__(\n        self,\n        next,\n        prompt,\n        state=None,\n        index=0,\n        mask=None,\n        end_token_id=None,\n    ):\n        batch_size, max_length = tf.shape(prompt)[0], tf.shape(prompt)[1]\n        # Make sure max length and start index are the same dtype.\n        index = tf.cast(index, max_length.dtype)\n\n        def create_beams(x):\n            \"\"\"Add initial beam state.\"\"\"\n            return tf.repeat(x, self.num_beams, axis=0)\n\n        def flatten_beams(x):\n            \"\"\"Combine the beam dim and batch dim.\"\"\"\n            flat_shape = [batch_size * self.num_beams] + x.shape.as_list()[2:]\n            return tf.reshape(x, shape=flat_shape)\n\n        def unflatten_beams(x):\n            \"\"\"Separate the beam dim and batch dim.\"\"\"\n            unflat_shape = [batch_size, self.num_beams] + x.shape.as_list()[1:]\n            return tf.reshape(x, shape=unflat_shape)\n\n        mask = tf.zeros_like(prompt, dtype=tf.bool) if mask is None else mask\n        # `tf.while_loop` will not accept `None` as a value for `loop_vars`.\n        state = () if state is None else state\n        # Add extra sequences for each beam.\n        prompt, mask = create_beams(prompt), create_beams(mask)\n        state = tf.nest.map_structure(create_beams, state)\n        # Setup the initial beam log-likelihoods.\n        # On the first loop, make sure only the original beam is considered.\n        log_probs = tf.constant([[0.0] + [-1e9] * (self.num_beams - 1)])\n        log_probs = flatten_beams(tf.repeat(log_probs, batch_size, axis=0))\n\n        def cond(prompt, state, index, log_probs):\n            if end_token_id is None:\n                return True\n            # Stop if all sequences have produced a *new* end_token_id.\n            end_tokens = (prompt == end_token_id) & (~mask)\n            prompt_done = tf.reduce_any(end_tokens, axis=-1)\n            return not tf.reduce_all(prompt_done)\n\n        def body(prompt, state, index, log_probs):\n            # Compute the softmax distribution for the next token.\n            logits, state = next(prompt, state, index)\n            vocab_size = tf.shape(logits)[-1]\n            probs = keras.activations.softmax(logits)\n\n            # Compute the running log-likelihood of each new candidate.\n            next_log_probs = tf.math.log(probs) + log_probs[..., tf.newaxis]\n            # Reshape `preds` to shape `(batch_size, num_beams * vocab_size)`.\n            next_log_probs = tf.reshape(next_log_probs, shape=[batch_size, -1])\n\n            # Compute the top beam indices and next tokens.\n            next_log_probs, indices = tf.math.top_k(\n                next_log_probs, k=self.num_beams, sorted=False\n            )\n            beam_indices = indices // vocab_size\n            next_token = flatten_beams(indices % vocab_size)\n            # We need `ensure_shape` as `top_k` will change the static shape.\n            next_log_probs = flatten_beams(next_log_probs)\n            log_probs = tf.ensure_shape(next_log_probs, log_probs.shape)\n\n            def gather_beams(x):\n                x = unflatten_beams(x)\n                x = tf.gather(x, beam_indices, axis=1, batch_dims=1)\n                return flatten_beams(x)\n\n            prompt = gather_beams(prompt)\n            state = tf.nest.map_structure(gather_beams, state)\n\n            # Update each beam with the next token.\n            next_token = tf.cast(next_token, prompt.dtype)\n            # Don't overwrite anywhere mask is True.\n            next_token = tf.where(mask[:, index], prompt[:, index], next_token)\n            # Update the prompt with the next token.\n            next_token = next_token[:, tf.newaxis]\n            prompt = dynamic_update_slice(prompt, next_token, [0, index])\n            # Return the iteration of the loop state.\n            return (prompt, state, index + 1, log_probs)\n\n        prompt, _, _, log_probs = tf.while_loop(\n            cond=cond,\n            body=body,\n            loop_vars=(prompt, state, index, log_probs),\n            maximum_iterations=(max_length - index),\n        )\n\n        # Gather the top beam at each batch index.\n        prompt, log_probs = unflatten_beams(prompt), unflatten_beams(log_probs)\n        top_beams = tf.math.argmax(log_probs, axis=-1)[:, tf.newaxis]\n        prompt = tf.gather(prompt, top_beams, axis=1, batch_dims=1)\n        return tf.squeeze(prompt, axis=1)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"num_beams\": self.num_beams,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/beam_sampler.py"}], "after_files": [{"content": "# Copyright 2023 The KerasNLP Authors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Beam Sampler.\"\"\"\n\nimport tensorflow as tf\nfrom tensorflow import keras\nfrom tensorflow.compiler.tf2xla.python.xla import dynamic_update_slice\n\nfrom keras_nlp.api_export import keras_nlp_export\nfrom keras_nlp.samplers.sampler import Sampler\nfrom keras_nlp.samplers.sampler import call_args_docstring\nfrom keras_nlp.utils.python_utils import format_docstring\n\n\n@format_docstring(call_args=call_args_docstring)\n@keras_nlp_export(\"keras_nlp.samplers.BeamSampler\")\nclass BeamSampler(Sampler):\n    \"\"\"Beam Sampler class.\n\n    This sampler implements beam search algorithm. At each time-step, beam\n    search keeps the beams (sequences) of the top `num_beams` highest\n    accumulated probabilities, and uses each one of the beams to predict\n    candidate next tokens.\n\n    Args:\n        num_beams: int. The number of beams that should be kept at each\n            time-step. `num_beams` should be strictly positive.\n        return_all_beams: bool. When set to `True`, the sampler will return the top prompt,\n            all prompts and their respective probabilities score.\n\n    Call Args:\n        {{call_args}}\n\n    Examples:\n    Return only the beam with the highest accumulated probability.\n    ```python\n    # Use a simple alphabet of lowercase characters to [0, 26).\n    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n    char_lookup = {v: k for k, v in int_lookup.items()}\n    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n\n    def next(prompt, state, index):\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((batch_size, vocab_size))\n        return logits, state\n\n    output = keras_nlp.samplers.BeamSampler()(\n        next=next,\n        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n        index=5,\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> \"zzzzzaaaaaaa\"\n    ```\n    Return all beams and their probabilities.\n    ```python\n    # Use a simple alphabet of lowercase characters to [0, 26).\n    int_lookup = {i: chr(i + ord('a')) for i in range(26)}\n    char_lookup = {v: k for k, v in int_lookup.items()}\n    batch_size, length, vocab_size = 1, 12, len(int_lookup)\n\n    def next(prompt, state, index):\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((batch_size, vocab_size))\n        return logits, state\n\n    output = keras_nlp.samplers.BeamSampler(return_all_beams=True)(\n        next=next,\n        prompt=tf.fill((batch_size, length,), char_lookup['z']),\n        index=5,\n    )\n\n    print(output[0].shape)\n    # >>> (1, 5, 12)\n    print(output[1].shape)\n    # >>> (1, 5)\n    print([\"\".join([int_lookup[i] for i in s]) for s in output[0][0].numpy()])\n    # >>> \"zzzzzaaaaaaa\"\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        num_beams=5,\n        return_all_beams=False,\n    ):\n        super().__init__()\n        self.num_beams = num_beams\n        self.return_all_beams = return_all_beams\n\n    def __call__(\n        self,\n        next,\n        prompt,\n        state=None,\n        index=0,\n        mask=None,\n        end_token_id=None,\n    ):\n        batch_size, max_length = tf.shape(prompt)[0], tf.shape(prompt)[1]\n        # Make sure max length and start index are the same dtype.\n        index = tf.cast(index, max_length.dtype)\n\n        def create_beams(x):\n            \"\"\"Add initial beam state.\"\"\"\n            return tf.repeat(x, self.num_beams, axis=0)\n\n        def flatten_beams(x):\n            \"\"\"Combine the beam dim and batch dim.\"\"\"\n            flat_shape = [batch_size * self.num_beams] + x.shape.as_list()[2:]\n            return tf.reshape(x, shape=flat_shape)\n\n        def unflatten_beams(x):\n            \"\"\"Separate the beam dim and batch dim.\"\"\"\n            unflat_shape = [batch_size, self.num_beams] + x.shape.as_list()[1:]\n            return tf.reshape(x, shape=unflat_shape)\n\n        mask = tf.zeros_like(prompt, dtype=tf.bool) if mask is None else mask\n        # `tf.while_loop` will not accept `None` as a value for `loop_vars`.\n        state = () if state is None else state\n        # Add extra sequences for each beam.\n        prompt, mask = create_beams(prompt), create_beams(mask)\n        state = tf.nest.map_structure(create_beams, state)\n        # Setup the initial beam log-likelihoods.\n        # On the first loop, make sure only the original beam is considered.\n        log_probs = tf.constant([[0.0] + [-1e9] * (self.num_beams - 1)])\n        log_probs = flatten_beams(tf.repeat(log_probs, batch_size, axis=0))\n\n        def cond(prompt, state, index, log_probs):\n            if end_token_id is None:\n                return True\n            # Stop if all sequences have produced a *new* end_token_id.\n            end_tokens = (prompt == end_token_id) & (~mask)\n            prompt_done = tf.reduce_any(end_tokens, axis=-1)\n            return not tf.reduce_all(prompt_done)\n\n        def body(prompt, state, index, log_probs):\n            # Compute the softmax distribution for the next token.\n            logits, state = next(prompt, state, index)\n            vocab_size = tf.shape(logits)[-1]\n            probs = keras.activations.softmax(logits)\n\n            # Compute the running log-likelihood of each new candidate.\n            next_log_probs = tf.math.log(probs) + log_probs[..., tf.newaxis]\n            # Reshape `preds` to shape `(batch_size, num_beams * vocab_size)`.\n            next_log_probs = tf.reshape(next_log_probs, shape=[batch_size, -1])\n\n            # Compute the top beam indices and next tokens.\n            next_log_probs, indices = tf.math.top_k(\n                next_log_probs, k=self.num_beams, sorted=False\n            )\n            beam_indices = indices // vocab_size\n            next_token = flatten_beams(indices % vocab_size)\n            # We need `ensure_shape` as `top_k` will change the static shape.\n            next_log_probs = flatten_beams(next_log_probs)\n            log_probs = tf.ensure_shape(next_log_probs, log_probs.shape)\n\n            def gather_beams(x):\n                x = unflatten_beams(x)\n                x = tf.gather(x, beam_indices, axis=1, batch_dims=1)\n                return flatten_beams(x)\n\n            prompt = gather_beams(prompt)\n            state = tf.nest.map_structure(gather_beams, state)\n\n            # Update each beam with the next token.\n            next_token = tf.cast(next_token, prompt.dtype)\n            # Don't overwrite anywhere mask is True.\n            next_token = tf.where(mask[:, index], prompt[:, index], next_token)\n            # Update the prompt with the next token.\n            next_token = next_token[:, tf.newaxis]\n            prompt = dynamic_update_slice(prompt, next_token, [0, index])\n            # Return the iteration of the loop state.\n            return (prompt, state, index + 1, log_probs)\n\n        prompt, _, _, log_probs = tf.while_loop(\n            cond=cond,\n            body=body,\n            loop_vars=(prompt, state, index, log_probs),\n            maximum_iterations=(max_length - index),\n        )\n\n        all_prompts = unflatten_beams(prompt)\n        all_log_probs = unflatten_beams(log_probs)\n\n        if self.return_all_beams:\n            sorted_indices = tf.argsort(\n                all_log_probs, axis=-1, direction=\"DESCENDING\"\n            )\n            sorted_log_probs = tf.gather(\n                all_log_probs, sorted_indices, axis=-1, batch_dims=1\n            )\n            sorted_prompts = tf.gather(\n                all_prompts, sorted_indices, axis=1, batch_dims=1\n            )\n            return sorted_prompts, sorted_log_probs\n        else:\n            # Gather the top beam at each batch index.\n            top_beams = tf.math.argmax(all_log_probs, axis=-1)[:, tf.newaxis]\n            prompt = tf.gather(all_prompts, top_beams, axis=1, batch_dims=1)\n            return tf.squeeze(prompt, axis=1)\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"num_beams\": self.num_beams,\n                \"return_all_beams\": self.return_all_beams,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/beam_sampler.py"}]}