Fix engine arg on GPT

PiperOrigin-RevId: 468576123
google-research · Aug 18, 2022 · 78a7adf · 78a7adf
1 parent 9dc6785
commit 78a7adf
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 551 deletions.
diff --git a/cascades/__init__.py b/cascades/__init__.py
@@ -16,7 +16,7 @@
 
 # A new PyPI release will be pushed everytime `__version__` is increased
 # When changing this, also update the CHANGELOG.md
-__version__ = '0.3.1'
+__version__ = '0.3.2'
 
 from cascades._src.distributions.base import UniformCategorical
 from cascades._src.distributions.gpt import GPT

diff --git a/cascades/_src/distributions/gpt.py b/cascades/_src/distributions/gpt.py
@@ -24,7 +24,9 @@
 import jax
 import openai
 
-openai.api_key = os.getenv('OPENAI_API_KEY')
+api_key = os.getenv('OPENAI_API_KEY')
+if api_key:
+ openai.api_key = api_key
 
 
 # TODO(ddohan): Persist cache to disk
@@ -42,7 +44,7 @@ class GPT(dists.Distribution):
 
  stop: Optional[Iterable[Text]] = ('\n',)
 
- engine = 'davinci-codex'
+ engine: str = 'davinci-codex'
  temperature: float = 0.7
  max_tokens: int = 128
  top_p: float = .95