fix: top_k_returns to top_k_return (#40)

AI21Labs · Jan 22, 2024 · fe3765c · fe3765c
1 parent 11dd4da
commit fe3765c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/ai21/clients/common/chat_base.py b/ai21/clients/common/chat_base.py
@@ -21,7 +21,7 @@ def create(
         max_tokens: Optional[int] = 300,
         min_tokens: Optional[int] = 0,
         top_p: Optional[float] = 1.0,
-        top_k_returns: Optional[int] = 0,
+        top_k_return: Optional[int] = 0,
         stop_sequences: Optional[List[str]] = None,
         frequency_penalty: Optional[Penalty] = None,
         presence_penalty: Optional[Penalty] = None,
@@ -43,7 +43,7 @@ def _create_body(
         max_tokens: Optional[int] = 300,
         min_tokens: Optional[int] = 0,
         top_p: Optional[float] = 1.0,
-        top_k_returns: Optional[int] = 0,
+        top_k_return: Optional[int] = 0,
         stop_sequences: Optional[List[str]] = None,
         frequency_penalty: Optional[Penalty] = None,
         presence_penalty: Optional[Penalty] = None,
@@ -58,7 +58,7 @@ def _create_body(
             "minTokens": min_tokens,
             "numResults": num_results,
             "topP": top_p,
-            "topKReturn": top_k_returns,
+            "topKReturn": top_k_return,
             "stopSequences": stop_sequences,
             "frequencyPenalty": None if frequency_penalty is None else frequency_penalty.to_dict(),
             "presencePenalty": None if presence_penalty is None else presence_penalty.to_dict(),

diff --git a/ai21/clients/studio/resources/studio_chat.py b/ai21/clients/studio/resources/studio_chat.py
@@ -19,7 +19,7 @@ def create(
         max_tokens: Optional[int] = 300,
         min_tokens: Optional[int] = 0,
         top_p: Optional[float] = 1.0,
-        top_k_returns: Optional[int] = 0,
+        top_k_return: Optional[int] = 0,
         stop_sequences: Optional[List[str]] = None,
         frequency_penalty: Optional[Penalty] = None,
         presence_penalty: Optional[Penalty] = None,
@@ -35,7 +35,7 @@ def create(
             max_tokens=max_tokens,
             min_tokens=min_tokens,
             top_p=top_p,
-            top_k_returns=top_k_returns,
+            top_k_return=top_k_return,
             stop_sequences=stop_sequences,
             frequency_penalty=frequency_penalty,
             presence_penalty=presence_penalty,