diff --git a/README.md b/README.md index 98ce3dff..98e1c244 100644 --- a/README.md +++ b/README.md @@ -149,7 +149,7 @@ messages = [ chat_completions = client.chat.completions.create( messages=messages, - model="jamba-instruct", + model="jamba-instruct-preview", ) ``` diff --git a/ai21/clients/studio/resources/chat/chat_completions.py b/ai21/clients/studio/resources/chat/chat_completions.py index 20ca4030..503bc65b 100644 --- a/ai21/clients/studio/resources/chat/chat_completions.py +++ b/ai21/clients/studio/resources/chat/chat_completions.py @@ -17,29 +17,19 @@ def create( self, model: str, messages: List[ChatMessage], - n: int | NotGiven = NOT_GIVEN, - logprobs: bool | NotGiven = NOT_GIVEN, - top_logprobs: int | NotGiven = NOT_GIVEN, max_tokens: int | NotGiven = NOT_GIVEN, temperature: float | NotGiven = NOT_GIVEN, top_p: float | NotGiven = NOT_GIVEN, stop: str | List[str] | NotGiven = NOT_GIVEN, - frequency_penalty: float | NotGiven = NOT_GIVEN, - presence_penalty: float | NotGiven = NOT_GIVEN, **kwargs: Any, ) -> ChatCompletionResponse: body = self._create_body( model=model, messages=messages, - n=n, - logprobs=logprobs, - top_logprobs=top_logprobs, stop=stop, temperature=temperature, max_tokens=max_tokens, top_p=top_p, - frequency_penalty=frequency_penalty, - presence_penalty=presence_penalty, **kwargs, ) @@ -51,15 +41,10 @@ def _create_body( self, model: str, messages: List[ChatMessage], - n: Optional[int] | NotGiven, - logprobs: Optional[bool] | NotGiven, - top_logprobs: Optional[int] | NotGiven, max_tokens: Optional[int] | NotGiven, temperature: Optional[float] | NotGiven, top_p: Optional[float] | NotGiven, stop: Optional[Union[str, List[str]]] | NotGiven, - frequency_penalty: Optional[float] | NotGiven, - presence_penalty: Optional[float] | NotGiven, **kwargs: Any, ) -> Dict[str, Any]: return remove_not_given( @@ -68,13 +53,8 @@ def _create_body( "messages": messages, "temperature": temperature, "maxTokens": max_tokens, - "n": n, "topP": top_p, - "logprobs": logprobs, - "topLogprobs": top_logprobs, "stop": stop, - "frequencyPenalty": frequency_penalty, - "presencePenalty": presence_penalty, **kwargs, } ) diff --git a/examples/studio/chat/chat_completions.py b/examples/studio/chat/chat_completions.py index 3a0eb903..e340a7e9 100644 --- a/examples/studio/chat/chat_completions.py +++ b/examples/studio/chat/chat_completions.py @@ -13,15 +13,10 @@ response = client.chat.completions.create( messages=messages, model="new-model-name", - n=2, - logprobs=True, - top_logprobs=2, max_tokens=100, temperature=0.7, top_p=1.0, stop=["\n"], - frequency_penalty=0.1, - presence_penalty=0.1, ) print(response) diff --git a/tests/integration_tests/clients/studio/test_chat_completions.py b/tests/integration_tests/clients/studio/test_chat_completions.py index a98f79eb..5aef4fff 100644 --- a/tests/integration_tests/clients/studio/test_chat_completions.py +++ b/tests/integration_tests/clients/studio/test_chat_completions.py @@ -6,7 +6,7 @@ from ai21.models.chat.chat_completion_response import ChatCompletionResponse -_MODEL = "new-model-name" +_MODEL = "jamba-instruct-preview" _MESSAGES = [ ChatMessage( content="Hello, I need help studying for the coming test, can you teach me about the US constitution? ", @@ -27,13 +27,9 @@ def test_chat_completion(): messages=messages, num_results=num_results, max_tokens=64, - logprobs=True, - top_logprobs=0.6, temperature=0.7, stop=["\n"], top_p=0.3, - frequency_penalty=0.2, - presence_penalty=0.4, ) assert isinstance(response, ChatCompletionResponse)