Update sample API to match what is used.

PiperOrigin-RevId: 644391496
google-research · Jun 18, 2024 · ad8bb7b · ad8bb7b
1 parent c366a7a
commit ad8bb7b
Show file tree

Hide file tree

Showing 38 changed files with 39 additions and 38 deletions.
diff --git a/cascades/__init__.py b/cascades/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/__init__.py b/cascades/_src/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/distributions/__init__.py b/cascades/_src/distributions/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/distributions/base.py b/cascades/_src/distributions/base.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/distributions/base_test.py b/cascades/_src/distributions/base_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/distributions/choose.py b/cascades/_src/distributions/choose.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/distributions/gpt.py b/cascades/_src/distributions/gpt.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/distributions/strings.py b/cascades/_src/distributions/strings.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/distributions/strings_test.py b/cascades/_src/distributions/strings_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/handlers.py b/cascades/_src/handlers.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/handlers_test.py b/cascades/_src/handlers_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/__init__.py b/cascades/_src/inference/__init__.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/base.py b/cascades/_src/inference/base.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/base_test.py b/cascades/_src/inference/base_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/beam_search.py b/cascades/_src/inference/beam_search.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/beam_search_test.py b/cascades/_src/inference/beam_search_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/dataset_sampler.py b/cascades/_src/inference/dataset_sampler.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/forward_sampling.py b/cascades/_src/inference/forward_sampling.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/forward_sampling_test.py b/cascades/_src/inference/forward_sampling_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/importance_sampling.py b/cascades/_src/inference/importance_sampling.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/importance_sampling_test.py b/cascades/_src/inference/importance_sampling_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/likelihood_weighting.py b/cascades/_src/inference/likelihood_weighting.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/likelihood_weighting_test.py b/cascades/_src/inference/likelihood_weighting_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/rejection_sampling.py b/cascades/_src/inference/rejection_sampling.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/inference/rejection_sampling_test.py b/cascades/_src/inference/rejection_sampling_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/interpreter.py b/cascades/_src/interpreter.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/interpreter_test.py b/cascades/_src/interpreter_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/sampler.py b/cascades/_src/sampler.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/_src/sampler_test.py b/cascades/_src/sampler_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/calculator.py b/cascades/examples/calculator.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/notebooks/trice.ipynb b/cascades/examples/notebooks/trice.ipynb
@@ -46,13 +46,14 @@
       "outputs": [],
       "source": [
         "# pylint:disable=unused-argument\n",
-        "def sample(params, context, num_steps, seed):\n",
+        "def sample(params, context, num_steps=256, temperature=1.0, *, seed):\n",
         "  \"\"\"Draws a sample continuation.\n",
         "\n",
         "  Args:\n",
         "    params: a PyTree of parameters.\n",
         "    context: a Python string or a list of Python strings.\n",
         "    num_steps: The maximum number of tokens to generate.\n",
+        "    temperature: The temperature to use.\n",
         "    seed: The random seed.\n",
         "\n",
         "  Returns:\n",

diff --git a/cascades/examples/scratchpad.py b/cascades/examples/scratchpad.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/scratchpad_test.py b/cascades/examples/scratchpad_test.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/tasks/bbh.py b/cascades/examples/tasks/bbh.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/tasks/constitutional_rl.py b/cascades/examples/tasks/constitutional_rl.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/tasks/gsm8k.py b/cascades/examples/tasks/gsm8k.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/tasks/math.py b/cascades/examples/tasks/math.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

diff --git a/cascades/examples/tasks/svamp.py b/cascades/examples/tasks/svamp.py
@@ -1,4 +1,4 @@
-# Copyright 2023 The cascades Authors.
+# Copyright 2024 The cascades Authors.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.