We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add a temperature argument to the Contrastive Sampler 
Similar to our other samplers, we should add the temperature argument to the `Contrastive Sampler` along with the necessary tests.
The approach for accomplishing the same can be seen in #951. 

I would like to take this up. Thanks!


--- 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/contrastive_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 """Contrastive 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.ContrastiveSampler")
28 class ContrastiveSampler(Sampler):
29     """Contrastive Sampler class.
30 
31     This sampler implements contrastive search algorithm. In short, the sampler
32     chooses the token having the max "score" as the next token. The "score" is
33     a weighted sum between token's probability and max similarity against
34     previous tokens. By using this joint score, contrastive sampler reduces the
35     behavior of duplicating seen tokens.
36 
37     Args:
38         k: int, the `k` value of top-k. Next token will be chosen from k tokens.
39         alpha: float, the weight of minus max similarity in joint score
40             computation. The larger the value of `alpha`, the score relies more
41             on the similarity than the token probability.
42         seed: int, defaults to None. The random seed.
43 
44     Call Args:
45         {{call_args}}
46 
47     Examples:
48     ```python
49     # Use a simple alphabet of lowercase characters to [0, 26).
50     int_lookup = {i: chr(i + ord("a")) for i in range(26)}
51     char_lookup = {v: k for k, v in int_lookup.items()}
52     batch_size, length, vocab_size = 1, 12, len(int_lookup)
53     hidden_size = 5
54     index = 5
55 
56     def next(prompt, cache, index):
57         prompt_batch_size = tf.shape(prompt)[0]
58         hidden_states = tf.ones((prompt_batch_size, hidden_size))
59         # A uniform distribution over our alphabet.
60         logits = tf.ones((prompt_batch_size, vocab_size))
61         return logits, hidden_states, cache
62 
63     output = keras_nlp.samplers.ContrastiveSampler()(
64         next=next,
65         prompt=tf.fill((batch_size, length), char_lookup["z"]),
66         index=index,
67         hidden_states=tf.ones([batch_size, index, hidden_size]),
68     )
69     print(["".join([int_lookup[i] for i in s]) for s in output.numpy()])
70     # >>> "zzzzzeeeeeee"
71     ```
72     """
73 
74     def __init__(
75         self,
76         k=5,
77         alpha=0.6,
78         seed=None,
79     ):
80         super().__init__()
81         self.k = k
82         self.alpha = alpha
83         self.seed = seed
84 
85     def __call__(
86         self,
87         next,
88         prompt,
89         cache=None,
90         index=0,
91         mask=None,
92         end_token_id=None,
93         hidden_states=None,
94     ):
95         if hidden_states is None:
96             raise ValueError(
97                 "`ContrastiveSampler` requires passing a `hidden_states`, but"
98                 "received `None`."
99             )
100         batch_size, max_length = tf.shape(prompt)[0], tf.shape(prompt)[1]
101         # Make sure max length and start index are the same dtype.
102         index = tf.cast(index, max_length.dtype)
103 
104         def create_beams(x):
105             """Add initial beam state."""
106             x = tf.repeat(x, self.k, axis=0)
107             flat_shape = [batch_size * self.k] + x.shape.as_list()[1:]
108             return tf.reshape(x, shape=flat_shape)
109 
110         def flatten_beams(x):
111             """Combine the beam dim and batch dim."""
112             flat_shape = [batch_size * self.k] + x.shape.as_list()[2:]
113             return tf.reshape(x, shape=flat_shape)
114 
115         def unflatten_beams(x):
116             """Separate the beam dim and batch dim."""
117             unflat_shape = [batch_size, self.k] + x.shape.as_list()[1:]
118             return tf.reshape(x, shape=unflat_shape)
119 
120         mask = tf.zeros_like(prompt, dtype=tf.bool) if mask is None else mask
121         # Compute initial logits.
122         logits, _, cache = next(prompt, cache, index)
123         # `tf.while_loop` will not accept `None` as a value for `loop_vars`.
124         cache = () if cache is None else cache
125 
126         def cond(prompt, cache, index, logits, hidden_states):
127             if end_token_id is None:
128                 return True
129             # Stop if all sequences have produced a *new* end_token_id.
130             end_tokens = (prompt == end_token_id) & (~mask)
131             prompt_done = tf.reduce_any(end_tokens, axis=-1)
132             return not tf.reduce_all(prompt_done)
133 
134         def body(prompt, cache, index, logits, hidden_states):
135             # Compute the softmax distribution for the next token.
136             probabilities = keras.activations.softmax(logits)
137 
138             # Replicate for `self.k` times to find the best token in top-k
139             # candidates.
140             prompt_beams = create_beams(prompt)
141             mask_beams = create_beams(mask)
142             hidden_states_beams = create_beams(hidden_states)
143             cache_beams = tf.nest.map_structure(create_beams, cache)
144 
145             # Get top-k candidate tokens and their probabilities.
146             top_k_probabilities, top_k_indices = tf.math.top_k(
147                 probabilities, k=self.k, sorted=False
148             )
149             next_token_probabilities = flatten_beams(top_k_probabilities)
150             next_token = flatten_beams(top_k_indices)
151             next_token = tf.cast(next_token, prompt.dtype)
152             next_token = tf.where(
153                 mask_beams[:, index], prompt_beams[:, index], next_token
154             )
155 
156             # Update the prompt with the next token.
157             next_token = next_token[:, tf.newaxis]
158             prompt_beams = dynamic_update_slice(
159                 prompt_beams, next_token, [0, index]
160             )
161 
162             # Compute the logits and hidden states for top-k candidate tokens.
163             next_logits, next_hidden_states_beams, cache_beams = next(
164                 prompt_beams, cache_beams, index + 1
165             )
166 
167             # Compute the max similarity score for top-k candidate tokens
168             # against previous tokens.
169             similarity_scores = self.similarity(
170                 hidden_states_beams, next_hidden_states_beams
171             )
172             max_similarity_scores = tf.cast(
173                 tf.reduce_max(similarity_scores[:, :index], axis=1),
174                 dtype=next_token_probabilities.dtype,
175             )
176             if index == 0:
177                 # If the index is 0, there is no previous states so we set
178                 # `max_similarity_scores` the same for all beams.
179                 max_similarity_scores = tf.zeros_like(max_similarity_scores)
180             # The final score of each candidate token is weighted sum of
181             # probability and similarity against previous tokens.
182             accumulated_scores = (
183                 (1 - self.alpha) * next_token_probabilities
184                 - self.alpha * max_similarity_scores
185             )
186             # Unflatten varibles to shape [batch_size, self.k, ...] for
187             # gather purpose.
188             unflat_score = unflatten_beams(accumulated_scores)
189             unflat_prompt = unflatten_beams(prompt_beams)
190             unflat_next_logits = unflatten_beams(next_logits)
191             unflat_next_hidden_states = unflatten_beams(
192                 next_hidden_states_beams
193             )
194             unflat_cache = tf.nest.map_structure(unflatten_beams, cache_beams)
195             best_token_indices = tf.math.argmax(unflat_score, axis=1)
196 
197             def gather_best_token(beams):
198                 return tf.gather(
199                     beams,
200                     best_token_indices,
201                     axis=1,
202                     batch_dims=1,
203                 )
204 
205             prompt = gather_best_token(unflat_prompt)
206             # We avoid recomputing forward pass for each token by updating the
207             # cache/hidden_states using the output, and pass the logits to
208             # next iteration step.
209             logits = gather_best_token(unflat_next_logits)
210             next_hidden_states = gather_best_token(unflat_next_hidden_states)
211             cache = tf.nest.map_structure(gather_best_token, unflat_cache)
212 
213             hidden_states = dynamic_update_slice(
214                 hidden_states,
215                 next_hidden_states[:, tf.newaxis, :],
216                 [0, index, 0],
217             )
218             return (prompt, cache, index + 1, logits, hidden_states)
219 
220         prompt, _, _, _, _ = tf.while_loop(
221             cond=cond,
222             body=body,
223             loop_vars=(prompt, cache, index, logits, hidden_states),
224             maximum_iterations=(max_length - index),
225         )
226         return prompt
227 
228     def similarity(self, h1, h2):
229         h2 = h2[..., tf.newaxis]
230         return tf.squeeze(tf.matmul(h1, h2), axis=-1) / (
231             tf.norm(h1, axis=-1) * tf.norm(h2, axis=-2)
232         )
233 
234     def get_config(self):
235         config = super().get_config()
236         config.update(
237             {
238                 "k": self.k,
239                 "alpha": self.alpha,
240                 "seed": self.seed,
241             }
242         )
243         return config
244
```
--- 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/contrastive_sampler.py b/keras_nlp/samplers/contrastive_sampler.py
--- a/keras_nlp/samplers/contrastive_sampler.py
+++ b/keras_nlp/samplers/contrastive_sampler.py
@@ -76,8 +76,9 @@
         k=5,
         alpha=0.6,
         seed=None,
+        **kwargs,
     ):
-        super().__init__()
+        super().__init__(**kwargs)
         self.k = k
         self.alpha = alpha
         self.seed = seed
@@ -133,7 +134,7 @@
 
         def body(prompt, cache, index, logits, hidden_states):
             # Compute the softmax distribution for the next token.
-            probabilities = keras.activations.softmax(logits)
+            probabilities = keras.activations.softmax(logits / self.temperature)
 
             # Replicate for `self.k` times to find the best token in top-k
             # candidates.

{"golden_diff": "diff --git a/keras_nlp/samplers/contrastive_sampler.py b/keras_nlp/samplers/contrastive_sampler.py\n--- a/keras_nlp/samplers/contrastive_sampler.py\n+++ b/keras_nlp/samplers/contrastive_sampler.py\n@@ -76,8 +76,9 @@\n         k=5,\n         alpha=0.6,\n         seed=None,\n+        **kwargs,\n     ):\n-        super().__init__()\n+        super().__init__(**kwargs)\n         self.k = k\n         self.alpha = alpha\n         self.seed = seed\n@@ -133,7 +134,7 @@\n \n         def body(prompt, cache, index, logits, hidden_states):\n             # Compute the softmax distribution for the next token.\n-            probabilities = keras.activations.softmax(logits)\n+            probabilities = keras.activations.softmax(logits / self.temperature)\n \n             # Replicate for `self.k` times to find the best token in top-k\n             # candidates.\n", "issue": "Add a temperature argument to the Contrastive Sampler \nSimilar to our other samplers, we should add the temperature argument to the `Contrastive Sampler` along with the necessary tests.\r\nThe approach for accomplishing the same can be seen in #951. \r\n\r\nI would like to take this up. Thanks!\r\n\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\"\"\"Contrastive 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.ContrastiveSampler\")\nclass ContrastiveSampler(Sampler):\n    \"\"\"Contrastive Sampler class.\n\n    This sampler implements contrastive search algorithm. In short, the sampler\n    chooses the token having the max \"score\" as the next token. The \"score\" is\n    a weighted sum between token's probability and max similarity against\n    previous tokens. By using this joint score, contrastive sampler reduces the\n    behavior of duplicating seen tokens.\n\n    Args:\n        k: int, the `k` value of top-k. Next token will be chosen from k tokens.\n        alpha: float, the weight of minus max similarity in joint score\n            computation. The larger the value of `alpha`, the score relies more\n            on the similarity than the token probability.\n        seed: int, defaults to None. The random seed.\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    hidden_size = 5\n    index = 5\n\n    def next(prompt, cache, index):\n        prompt_batch_size = tf.shape(prompt)[0]\n        hidden_states = tf.ones((prompt_batch_size, hidden_size))\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((prompt_batch_size, vocab_size))\n        return logits, hidden_states, cache\n\n    output = keras_nlp.samplers.ContrastiveSampler()(\n        next=next,\n        prompt=tf.fill((batch_size, length), char_lookup[\"z\"]),\n        index=index,\n        hidden_states=tf.ones([batch_size, index, hidden_size]),\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> \"zzzzzeeeeeee\"\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        k=5,\n        alpha=0.6,\n        seed=None,\n    ):\n        super().__init__()\n        self.k = k\n        self.alpha = alpha\n        self.seed = seed\n\n    def __call__(\n        self,\n        next,\n        prompt,\n        cache=None,\n        index=0,\n        mask=None,\n        end_token_id=None,\n        hidden_states=None,\n    ):\n        if hidden_states is None:\n            raise ValueError(\n                \"`ContrastiveSampler` requires passing a `hidden_states`, but\"\n                \"received `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            x = tf.repeat(x, self.k, axis=0)\n            flat_shape = [batch_size * self.k] + x.shape.as_list()[1:]\n            return tf.reshape(x, shape=flat_shape)\n\n        def flatten_beams(x):\n            \"\"\"Combine the beam dim and batch dim.\"\"\"\n            flat_shape = [batch_size * self.k] + 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.k] + 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        # Compute initial logits.\n        logits, _, cache = next(prompt, cache, index)\n        # `tf.while_loop` will not accept `None` as a value for `loop_vars`.\n        cache = () if cache is None else cache\n\n        def cond(prompt, cache, index, logits, hidden_states):\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, cache, index, logits, hidden_states):\n            # Compute the softmax distribution for the next token.\n            probabilities = keras.activations.softmax(logits)\n\n            # Replicate for `self.k` times to find the best token in top-k\n            # candidates.\n            prompt_beams = create_beams(prompt)\n            mask_beams = create_beams(mask)\n            hidden_states_beams = create_beams(hidden_states)\n            cache_beams = tf.nest.map_structure(create_beams, cache)\n\n            # Get top-k candidate tokens and their probabilities.\n            top_k_probabilities, top_k_indices = tf.math.top_k(\n                probabilities, k=self.k, sorted=False\n            )\n            next_token_probabilities = flatten_beams(top_k_probabilities)\n            next_token = flatten_beams(top_k_indices)\n            next_token = tf.cast(next_token, prompt.dtype)\n            next_token = tf.where(\n                mask_beams[:, index], prompt_beams[:, index], next_token\n            )\n\n            # Update the prompt with the next token.\n            next_token = next_token[:, tf.newaxis]\n            prompt_beams = dynamic_update_slice(\n                prompt_beams, next_token, [0, index]\n            )\n\n            # Compute the logits and hidden states for top-k candidate tokens.\n            next_logits, next_hidden_states_beams, cache_beams = next(\n                prompt_beams, cache_beams, index + 1\n            )\n\n            # Compute the max similarity score for top-k candidate tokens\n            # against previous tokens.\n            similarity_scores = self.similarity(\n                hidden_states_beams, next_hidden_states_beams\n            )\n            max_similarity_scores = tf.cast(\n                tf.reduce_max(similarity_scores[:, :index], axis=1),\n                dtype=next_token_probabilities.dtype,\n            )\n            if index == 0:\n                # If the index is 0, there is no previous states so we set\n                # `max_similarity_scores` the same for all beams.\n                max_similarity_scores = tf.zeros_like(max_similarity_scores)\n            # The final score of each candidate token is weighted sum of\n            # probability and similarity against previous tokens.\n            accumulated_scores = (\n                (1 - self.alpha) * next_token_probabilities\n                - self.alpha * max_similarity_scores\n            )\n            # Unflatten varibles to shape [batch_size, self.k, ...] for\n            # gather purpose.\n            unflat_score = unflatten_beams(accumulated_scores)\n            unflat_prompt = unflatten_beams(prompt_beams)\n            unflat_next_logits = unflatten_beams(next_logits)\n            unflat_next_hidden_states = unflatten_beams(\n                next_hidden_states_beams\n            )\n            unflat_cache = tf.nest.map_structure(unflatten_beams, cache_beams)\n            best_token_indices = tf.math.argmax(unflat_score, axis=1)\n\n            def gather_best_token(beams):\n                return tf.gather(\n                    beams,\n                    best_token_indices,\n                    axis=1,\n                    batch_dims=1,\n                )\n\n            prompt = gather_best_token(unflat_prompt)\n            # We avoid recomputing forward pass for each token by updating the\n            # cache/hidden_states using the output, and pass the logits to\n            # next iteration step.\n            logits = gather_best_token(unflat_next_logits)\n            next_hidden_states = gather_best_token(unflat_next_hidden_states)\n            cache = tf.nest.map_structure(gather_best_token, unflat_cache)\n\n            hidden_states = dynamic_update_slice(\n                hidden_states,\n                next_hidden_states[:, tf.newaxis, :],\n                [0, index, 0],\n            )\n            return (prompt, cache, index + 1, logits, hidden_states)\n\n        prompt, _, _, _, _ = tf.while_loop(\n            cond=cond,\n            body=body,\n            loop_vars=(prompt, cache, index, logits, hidden_states),\n            maximum_iterations=(max_length - index),\n        )\n        return prompt\n\n    def similarity(self, h1, h2):\n        h2 = h2[..., tf.newaxis]\n        return tf.squeeze(tf.matmul(h1, h2), axis=-1) / (\n            tf.norm(h1, axis=-1) * tf.norm(h2, axis=-2)\n        )\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"k\": self.k,\n                \"alpha\": self.alpha,\n                \"seed\": self.seed,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/contrastive_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\"\"\"Contrastive 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.ContrastiveSampler\")\nclass ContrastiveSampler(Sampler):\n    \"\"\"Contrastive Sampler class.\n\n    This sampler implements contrastive search algorithm. In short, the sampler\n    chooses the token having the max \"score\" as the next token. The \"score\" is\n    a weighted sum between token's probability and max similarity against\n    previous tokens. By using this joint score, contrastive sampler reduces the\n    behavior of duplicating seen tokens.\n\n    Args:\n        k: int, the `k` value of top-k. Next token will be chosen from k tokens.\n        alpha: float, the weight of minus max similarity in joint score\n            computation. The larger the value of `alpha`, the score relies more\n            on the similarity than the token probability.\n        seed: int, defaults to None. The random seed.\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    hidden_size = 5\n    index = 5\n\n    def next(prompt, cache, index):\n        prompt_batch_size = tf.shape(prompt)[0]\n        hidden_states = tf.ones((prompt_batch_size, hidden_size))\n        # A uniform distribution over our alphabet.\n        logits = tf.ones((prompt_batch_size, vocab_size))\n        return logits, hidden_states, cache\n\n    output = keras_nlp.samplers.ContrastiveSampler()(\n        next=next,\n        prompt=tf.fill((batch_size, length), char_lookup[\"z\"]),\n        index=index,\n        hidden_states=tf.ones([batch_size, index, hidden_size]),\n    )\n    print([\"\".join([int_lookup[i] for i in s]) for s in output.numpy()])\n    # >>> \"zzzzzeeeeeee\"\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        k=5,\n        alpha=0.6,\n        seed=None,\n        **kwargs,\n    ):\n        super().__init__(**kwargs)\n        self.k = k\n        self.alpha = alpha\n        self.seed = seed\n\n    def __call__(\n        self,\n        next,\n        prompt,\n        cache=None,\n        index=0,\n        mask=None,\n        end_token_id=None,\n        hidden_states=None,\n    ):\n        if hidden_states is None:\n            raise ValueError(\n                \"`ContrastiveSampler` requires passing a `hidden_states`, but\"\n                \"received `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            x = tf.repeat(x, self.k, axis=0)\n            flat_shape = [batch_size * self.k] + x.shape.as_list()[1:]\n            return tf.reshape(x, shape=flat_shape)\n\n        def flatten_beams(x):\n            \"\"\"Combine the beam dim and batch dim.\"\"\"\n            flat_shape = [batch_size * self.k] + 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.k] + 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        # Compute initial logits.\n        logits, _, cache = next(prompt, cache, index)\n        # `tf.while_loop` will not accept `None` as a value for `loop_vars`.\n        cache = () if cache is None else cache\n\n        def cond(prompt, cache, index, logits, hidden_states):\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, cache, index, logits, hidden_states):\n            # Compute the softmax distribution for the next token.\n            probabilities = keras.activations.softmax(logits / self.temperature)\n\n            # Replicate for `self.k` times to find the best token in top-k\n            # candidates.\n            prompt_beams = create_beams(prompt)\n            mask_beams = create_beams(mask)\n            hidden_states_beams = create_beams(hidden_states)\n            cache_beams = tf.nest.map_structure(create_beams, cache)\n\n            # Get top-k candidate tokens and their probabilities.\n            top_k_probabilities, top_k_indices = tf.math.top_k(\n                probabilities, k=self.k, sorted=False\n            )\n            next_token_probabilities = flatten_beams(top_k_probabilities)\n            next_token = flatten_beams(top_k_indices)\n            next_token = tf.cast(next_token, prompt.dtype)\n            next_token = tf.where(\n                mask_beams[:, index], prompt_beams[:, index], next_token\n            )\n\n            # Update the prompt with the next token.\n            next_token = next_token[:, tf.newaxis]\n            prompt_beams = dynamic_update_slice(\n                prompt_beams, next_token, [0, index]\n            )\n\n            # Compute the logits and hidden states for top-k candidate tokens.\n            next_logits, next_hidden_states_beams, cache_beams = next(\n                prompt_beams, cache_beams, index + 1\n            )\n\n            # Compute the max similarity score for top-k candidate tokens\n            # against previous tokens.\n            similarity_scores = self.similarity(\n                hidden_states_beams, next_hidden_states_beams\n            )\n            max_similarity_scores = tf.cast(\n                tf.reduce_max(similarity_scores[:, :index], axis=1),\n                dtype=next_token_probabilities.dtype,\n            )\n            if index == 0:\n                # If the index is 0, there is no previous states so we set\n                # `max_similarity_scores` the same for all beams.\n                max_similarity_scores = tf.zeros_like(max_similarity_scores)\n            # The final score of each candidate token is weighted sum of\n            # probability and similarity against previous tokens.\n            accumulated_scores = (\n                (1 - self.alpha) * next_token_probabilities\n                - self.alpha * max_similarity_scores\n            )\n            # Unflatten varibles to shape [batch_size, self.k, ...] for\n            # gather purpose.\n            unflat_score = unflatten_beams(accumulated_scores)\n            unflat_prompt = unflatten_beams(prompt_beams)\n            unflat_next_logits = unflatten_beams(next_logits)\n            unflat_next_hidden_states = unflatten_beams(\n                next_hidden_states_beams\n            )\n            unflat_cache = tf.nest.map_structure(unflatten_beams, cache_beams)\n            best_token_indices = tf.math.argmax(unflat_score, axis=1)\n\n            def gather_best_token(beams):\n                return tf.gather(\n                    beams,\n                    best_token_indices,\n                    axis=1,\n                    batch_dims=1,\n                )\n\n            prompt = gather_best_token(unflat_prompt)\n            # We avoid recomputing forward pass for each token by updating the\n            # cache/hidden_states using the output, and pass the logits to\n            # next iteration step.\n            logits = gather_best_token(unflat_next_logits)\n            next_hidden_states = gather_best_token(unflat_next_hidden_states)\n            cache = tf.nest.map_structure(gather_best_token, unflat_cache)\n\n            hidden_states = dynamic_update_slice(\n                hidden_states,\n                next_hidden_states[:, tf.newaxis, :],\n                [0, index, 0],\n            )\n            return (prompt, cache, index + 1, logits, hidden_states)\n\n        prompt, _, _, _, _ = tf.while_loop(\n            cond=cond,\n            body=body,\n            loop_vars=(prompt, cache, index, logits, hidden_states),\n            maximum_iterations=(max_length - index),\n        )\n        return prompt\n\n    def similarity(self, h1, h2):\n        h2 = h2[..., tf.newaxis]\n        return tf.squeeze(tf.matmul(h1, h2), axis=-1) / (\n            tf.norm(h1, axis=-1) * tf.norm(h2, axis=-2)\n        )\n\n    def get_config(self):\n        config = super().get_config()\n        config.update(\n            {\n                \"k\": self.k,\n                \"alpha\": self.alpha,\n                \"seed\": self.seed,\n            }\n        )\n        return config\n", "path": "keras_nlp/samplers/contrastive_sampler.py"}]}