feat: Sunset tsms and j2

.gitignore

@@ -119,3 +119,5 @@ fabric.properties
 *.orig
 *.rej
 __pycache__
+
+tests/integration_tests/test-file/*

README.md

@@ -304,7 +304,7 @@ Note that jamba-instruct supports async and streaming as well.
 
 </details>
 
-For a more detailed example, see the completion [examples](examples/studio/completion.py).
+For a more detailed example, see the completion [examples](examples/studio/chat/chat_completions.py).
 
 ---
 
@@ -388,33 +388,8 @@ For a more detailed example, see the chat [sync](examples/studio/conversational_
 
 ---
 
-## More Models
-
-## TSMs
-
-AI21 Studio's Task-Specific Models offer a range of powerful tools. These models have been specifically designed for their respective tasks and provide high-quality results while optimizing efficiency.
-The full documentation and guides can be found [here](https://docs.ai21.com/docs/task-specific).
-
-### Contextual Answers
-
-The `answer` API allows you to access our high-quality question answering model.
-
-```python
-from ai21 import AI21Client
-
-client = AI21Client()
-response = client.answer.create(
-    context="This is a text is for testing purposes",
-    question="Question about context",
-)
-```
-
-A detailed explanation on Contextual Answers, can be found [here](https://docs.ai21.com/docs/contextual-answers-api)
-
 ### File Upload
 
----
-
 ```python
 from ai21 import AI21Client
 
@@ -430,8 +405,6 @@ file_id = client.library.files.create(
 uploaded_file = client.library.files.get(file_id)
 ```
 
-For more information on more Task Specific Models, see the [documentation](https://docs.ai21.com/reference/paraphrase-api-ref).
-
 ## Token Counting
 
 ---

ai21/__init__.py

@@ -14,7 +14,6 @@
     TooManyRequestsError,
 )
 from ai21.logger import setup_logger
-from ai21.services.sagemaker import SageMaker
 from ai21.version import VERSION
 
 __version__ = VERSION
@@ -27,12 +26,6 @@ def _import_bedrock_client():
     return AI21BedrockClient
 
 
-def _import_sagemaker_client():
-    from ai21.clients.sagemaker.ai21_sagemaker_client import AI21SageMakerClient
-
-    return AI21SageMakerClient
-
-
 def _import_bedrock_model_id():
     from ai21.clients.bedrock.bedrock_model_id import BedrockModelID
 
@@ -45,12 +38,6 @@ def _import_async_bedrock_client():
     return AsyncAI21BedrockClient
 
 
-def _import_async_sagemaker_client():
-    from ai21.clients.sagemaker.ai21_sagemaker_client import AsyncAI21SageMakerClient
-
-    return AsyncAI21SageMakerClient
-
-
 def _import_vertex_client():
     from ai21.clients.vertex.ai21_vertex_client import AI21VertexClient
 
@@ -68,25 +55,19 @@ def __getattr__(name: str) -> Any:
         if name == "AI21BedrockClient":
             return _import_bedrock_client()
 
-        if name == "AI21SageMakerClient":
-            return _import_sagemaker_client()
-
         if name == "BedrockModelID":
             return _import_bedrock_model_id()
 
         if name == "AsyncAI21BedrockClient":
             return _import_async_bedrock_client()
 
-        if name == "AsyncAI21SageMakerClient":
-            return _import_async_sagemaker_client()
-
         if name == "AI21VertexClient":
             return _import_vertex_client()
 
         if name == "AsyncAI21VertexClient":
             return _import_async_vertex_client()
     except ImportError as e:
-        raise ImportError('Please install "ai21[AWS]" for SageMaker or Bedrock, or "ai21[Vertex]" for Vertex') from e
+        raise ImportError('Please install "ai21[AWS]" for Bedrock, or "ai21[Vertex]" for Vertex') from e
 
 
 __all__ = [
@@ -100,13 +81,10 @@ def __getattr__(name: str) -> Any:
     "ModelPackageDoesntExistError",
     "TooManyRequestsError",
     "AI21BedrockClient",
-    "AI21SageMakerClient",
     "BedrockModelID",
-    "SageMaker",
     "AI21AzureClient",
     "AsyncAI21AzureClient",
     "AsyncAI21BedrockClient",
-    "AsyncAI21SageMakerClient",
     "AI21VertexClient",
     "AsyncAI21VertexClient",
 ]

ai21/clients/common/completion_base.py

@@ -42,7 +42,6 @@ def create(
         temperature: float | NOT_GIVEN = NOT_GIVEN,
         top_p: float | NotGiven = NOT_GIVEN,
         top_k_return: int | NotGiven = NOT_GIVEN,
-        custom_model: str | NotGiven = NOT_GIVEN,
         stop_sequences: List[str] | NotGiven = NOT_GIVEN,
         frequency_penalty: Penalty | NotGiven = NOT_GIVEN,
         presence_penalty: Penalty | NotGiven = NOT_GIVEN,
@@ -60,7 +59,6 @@ def create(
         :param temperature: A value controlling the "creativity" of the model's responses.
         :param top_p: A value controlling the diversity of the model's responses.
         :param top_k_return: The number of top-scoring tokens to consider for each generation step.
-        :param custom_model:
         :param stop_sequences: Stops decoding if any of the strings is generated
         :param frequency_penalty: A penalty applied to tokens that are frequently generated.
         :param presence_penalty:  A penalty applied to tokens that are already present in the prompt.
@@ -84,7 +82,6 @@ def _create_body(
         temperature: float | NotGiven,
         top_p: float | NotGiven,
         top_k_return: int | NotGiven,
-        custom_model: str | NotGiven,
         stop_sequences: List[str] | NotGiven,
         frequency_penalty: Penalty | NotGiven,
         presence_penalty: Penalty | NotGiven,
@@ -96,7 +93,6 @@ def _create_body(
         return remove_not_given(
             {
                 "model": model,
-                "customModel": custom_model,
                 "prompt": prompt,
                 "maxTokens": max_tokens,
                 "numResults": num_results,
@@ -114,12 +110,5 @@ def _create_body(
             }
         )
 
-    def _get_completion_path(self, model: str, custom_model: str | NotGiven = NOT_GIVEN):
-        path = f"/{model}"
-
-        if custom_model:
-            path = f"{path}/{custom_model}"
-
-        path = f"{path}/{self._module_name}"
-
-        return path
+    def _get_completion_path(self, model: str):
+        return f"/{model}/{self._module_name}"