snowflake-ml-python 1.10.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

snowflake/ml/model/_client/service/model_deployment_spec_schema.py

@@ -10,10 +10,15 @@ class Model(BaseModel):
     version: str
 
 
+class InferenceEngineSpec(BaseModel):
+    inference_engine_name: str
+    inference_engine_args: Optional[list[str]] = None
+
+
 class ImageBuild(BaseModel):
-    compute_pool: str
-    image_repo: str
-    force_rebuild: bool
+    compute_pool: Optional[str] = None
+    image_repo: Optional[str] = None
+    force_rebuild: Optional[bool] = None
     external_access_integrations: Optional[list[str]] = None
 
 
@@ -27,6 +32,17 @@ class Service(BaseModel):
     gpu: Optional[str] = None
     num_workers: Optional[int] = None
     max_batch_rows: Optional[int] = None
+    inference_engine_spec: Optional[InferenceEngineSpec] = None
+
+
+class Input(BaseModel):
+    input_stage_location: str
+    input_file_pattern: str
+
+
+class Output(BaseModel):
+    output_stage_location: str
+    completion_filename: str
 
 
 class Job(BaseModel):
@@ -37,10 +53,10 @@ class Job(BaseModel):
     gpu: Optional[str] = None
     num_workers: Optional[int] = None
     max_batch_rows: Optional[int] = None
-    warehouse: str
-    target_method: str
-    input_table_name: str
-    output_table_name: str
+    warehouse: Optional[str] = None
+    function_name: str
+    input: Input
+    output: Output
 
 
 class LogModelArgs(BaseModel):
@@ -68,13 +84,13 @@ class ModelLogging(BaseModel):
 
 class ModelServiceDeploymentSpec(BaseModel):
     models: list[Model]
-    image_build: ImageBuild
+    image_build: Optional[ImageBuild] = None
     service: Service
     model_loggings: Optional[list[ModelLogging]] = None
 
 
 class ModelJobDeploymentSpec(BaseModel):
     models: list[Model]
-    image_build: ImageBuild
+    image_build: Optional[ImageBuild] = None
     job: Job
     model_loggings: Optional[list[ModelLogging]] = None

snowflake/ml/model/_model_composer/model_composer.py

@@ -1,17 +1,12 @@
 import pathlib
 import tempfile
 import uuid
-import warnings
 from types import ModuleType
 from typing import TYPE_CHECKING, Any, Optional, Union
 from urllib import parse
 
-from absl import logging
-from packaging import requirements
-
 from snowflake import snowpark
-from snowflake.ml import version as snowml_version
-from snowflake.ml._internal import env as snowml_env, env_utils, file_utils
+from snowflake.ml._internal import file_utils
 from snowflake.ml._internal.lineage import lineage_utils
 from snowflake.ml.data import data_source
 from snowflake.ml.model import model_signature, type_hints as model_types
@@ -19,7 +14,6 @@ from snowflake.ml.model._model_composer.model_manifest import model_manifest
 from snowflake.ml.model._packager import model_packager
 from snowflake.ml.model._packager.model_meta import model_meta
 from snowflake.snowpark import Session
-from snowflake.snowpark._internal import utils as snowpark_utils
 
 if TYPE_CHECKING:
     from snowflake.ml.experiment._experiment_info import ExperimentInfo
@@ -142,73 +136,10 @@ class ModelComposer:
         experiment_info: Optional["ExperimentInfo"] = None,
         options: Optional[model_types.ModelSaveOption] = None,
     ) -> model_meta.ModelMetadata:
-        # set enable_explainability=False if the model is not runnable in WH or the target platforms include SPCS
-        conda_dep_dict = env_utils.validate_conda_dependency_string_list(
-            conda_dependencies if conda_dependencies else []
-        )
-
-        enable_explainability = None
-
-        if options:
-            enable_explainability = options.get("enable_explainability", None)
-
-        # skip everything if user said False explicitly
-        if enable_explainability is None or enable_explainability is True:
-            is_warehouse_runnable = (
-                not conda_dep_dict
-                or all(
-                    chan == env_utils.DEFAULT_CHANNEL_NAME or chan == env_utils.SNOWFLAKE_CONDA_CHANNEL_URL
-                    for chan in conda_dep_dict
-                )
-            ) and (not pip_requirements)
-
-            only_spcs = (
-                target_platforms
-                and len(target_platforms) == 1
-                and model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES in target_platforms
-            )
-            if only_spcs or (not is_warehouse_runnable):
-                # if only SPCS and user asked for explainability we fail
-                if enable_explainability is True:
-                    raise ValueError(
-                        "`enable_explainability` cannot be set to True when the model is not runnable in WH "
-                        "or the target platforms include SPCS."
-                    )
-                elif not options:  # explicitly set flag to false in these cases if not specified
-                    options = model_types.BaseModelSaveOption()
-                    options["enable_explainability"] = False
-            elif (
-                target_platforms
-                and len(target_platforms) > 1
-                and model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES in target_platforms
-            ):  # if both then only available for WH
-                if enable_explainability is True:
-                    warnings.warn(
-                        ("Explain function will only be available for model deployed to warehouse."),
-                        category=UserWarning,
-                        stacklevel=2,
-                    )
 
         if not options:
             options = model_types.BaseModelSaveOption()
 
-        if not snowpark_utils.is_in_stored_procedure() and target_platforms != [  # type: ignore[no-untyped-call]
-            model_types.TargetPlatform.SNOWPARK_CONTAINER_SERVICES  # no information schema check for SPCS-only models
-        ]:
-            snowml_matched_versions = env_utils.get_matched_package_versions_in_information_schema(
-                self.session,
-                reqs=[requirements.Requirement(f"{env_utils.SNOWPARK_ML_PKG_NAME}=={snowml_version.VERSION}")],
-                python_version=python_version or snowml_env.PYTHON_VERSION,
-                statement_params=self._statement_params,
-            ).get(env_utils.SNOWPARK_ML_PKG_NAME, [])
-
-            if len(snowml_matched_versions) < 1 and options.get("embed_local_ml_library", False) is False:
-                logging.info(
-                    f"Local snowflake-ml-python library has version {snowml_version.VERSION},"
-                    " which is not available in the Snowflake server, embedding local ML library automatically."
-                )
-                options["embed_local_ml_library"] = True
-
         model_metadata: model_meta.ModelMetadata = self.packager.save(
             name=name,
             model=model,

snowflake/ml/model/_model_composer/model_manifest/model_manifest.py

@@ -1,13 +1,11 @@
 import collections
 import logging
 import pathlib
-import warnings
 from typing import TYPE_CHECKING, Optional, cast
 
 import yaml
 
 from snowflake.ml._internal import env_utils
-from snowflake.ml._internal.exceptions import error_codes, exceptions
 from snowflake.ml.data import data_source
 from snowflake.ml.model import type_hints
 from snowflake.ml.model._model_composer.model_manifest import model_manifest_schema
@@ -55,47 +53,8 @@ class ModelManifest:
         experiment_info: Optional["ExperimentInfo"] = None,
         target_platforms: Optional[list[type_hints.TargetPlatform]] = None,
     ) -> None:
-        if options is None:
-            options = {}
-
-        has_pip_requirements = len(model_meta.env.pip_requirements) > 0
-        only_spcs = (
-            target_platforms
-            and len(target_platforms) == 1
-            and target_platforms[0] == type_hints.TargetPlatform.SNOWPARK_CONTAINER_SERVICES
-        )
-
-        if "relax_version" not in options:
-            if has_pip_requirements or only_spcs:
-                logger.info(
-                    "Setting `relax_version=False` as this model will run in Snowpark Container Services "
-                    "or in Warehouse with a specified artifact_repository_map where exact version "
-                    " specifications will be honored."
-                )
-                relax_version = False
-            else:
-                warnings.warn(
-                    (
-                        "`relax_version` is not set and therefore defaulted to True. Dependency version constraints"
-                        " relaxed from ==x.y.z to >=x.y, <(x+1). To use specific dependency versions for compatibility,"
-                        " reproducibility, etc., set `options={'relax_version': False}` when logging the model."
-                    ),
-                    category=UserWarning,
-                    stacklevel=2,
-                )
-                relax_version = True
-            options["relax_version"] = relax_version
-        else:
-            relax_version = options.get("relax_version", True)
-            if relax_version and (has_pip_requirements or only_spcs):
-                raise exceptions.SnowflakeMLException(
-                    error_code=error_codes.INVALID_ARGUMENT,
-                    original_exception=ValueError(
-                        "Setting `relax_version=True` is only allowed for models to be run in Warehouse with "
-                        "Snowflake Conda Channel dependencies. It cannot be used with pip requirements or when "
-                        "targeting only Snowpark Container Services."
-                    ),
-                )
+        assert options is not None, "ModelParameterReconciler should have set options with relax_version"
+        relax_version = options["relax_version"]
 
         runtime_to_use = model_runtime.ModelRuntime(
             name=self._DEFAULT_RUNTIME_NAME,

snowflake/ml/model/_packager/model_handlers/huggingface_pipeline.py

@@ -1,6 +1,8 @@
 import json
 import logging
 import os
+import time
+import uuid
 import warnings
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
@@ -11,7 +13,12 @@ from packaging import version
 from typing_extensions import TypeGuard, Unpack
 
 from snowflake.ml._internal import type_utils
-from snowflake.ml.model import custom_model, model_signature, type_hints as model_types
+from snowflake.ml.model import (
+    custom_model,
+    model_signature,
+    openai_signatures,
+    type_hints as model_types,
+)
 from snowflake.ml.model._packager.model_env import model_env
 from snowflake.ml.model._packager.model_handlers import _base, _utils as handlers_utils
 from snowflake.ml.model._packager.model_handlers_migrator import base_migrator
@@ -151,7 +158,10 @@ class HuggingFacePipelineHandler(
             assert isinstance(model, huggingface_pipeline.HuggingFacePipelineModel)
             params = {**model.__dict__, **model.model_kwargs}
 
-        inferred_pipe_sig = model_signature_utils.huggingface_pipeline_signature_auto_infer(task, params=params)
+        inferred_pipe_sig = model_signature_utils.huggingface_pipeline_signature_auto_infer(
+            task,
+            params=params,
+        )
 
         if not is_sub_model:
             target_methods = handlers_utils.get_target_methods(
@@ -401,6 +411,34 @@ class HuggingFacePipelineHandler(
                             ),
                             axis=1,
                         ).to_list()
+                    elif raw_model.task == "text-generation":
+                        # verify when the target method is __call__ and
+                        # if the signature is default text-generation signature
+                        # then use the HuggingFaceOpenAICompatibleModel to wrap the pipeline
+                        if signature == openai_signatures._OPENAI_CHAT_SIGNATURE_SPEC:
+                            wrapped_model = HuggingFaceOpenAICompatibleModel(pipeline=raw_model)
+
+                            temp_res = X.apply(
+                                lambda row: wrapped_model.generate_chat_completion(
+                                    messages=row["messages"],
+                                    max_completion_tokens=row.get("max_completion_tokens", None),
+                                    temperature=row.get("temperature", None),
+                                    stop_strings=row.get("stop", None),
+                                    n=row.get("n", 1),
+                                    stream=row.get("stream", False),
+                                    top_p=row.get("top_p", 1.0),
+                                    frequency_penalty=row.get("frequency_penalty", None),
+                                    presence_penalty=row.get("presence_penalty", None),
+                                ),
+                                axis=1,
+                            ).to_list()
+                        else:
+                            if len(signature.inputs) > 1:
+                                input_data = X.to_dict("records")
+                            # If it is only expecting one argument, Then it is expecting a list of something.
+                            else:
+                                input_data = X[signature.inputs[0].name].to_list()
+                            temp_res = getattr(raw_model, target_method)(input_data)
                     else:
                         # For others, we could offer the whole dataframe as a list.
                         # Some of them may need some conversion
@@ -527,3 +565,170 @@ class HuggingFacePipelineHandler(
         hg_pipe_model = _HFPipelineModel(custom_model.ModelContext())
 
         return hg_pipe_model
+
+
+class HuggingFaceOpenAICompatibleModel:
+    """
+    A class to wrap a Hugging Face text generation model and provide an
+    OpenAI-compatible chat completion interface.
+    """
+
+    def __init__(self, pipeline: "transformers.Pipeline") -> None:
+        """
+        Initializes the model and tokenizer.
+
+        Args:
+            pipeline (transformers.pipeline): The Hugging Face pipeline to wrap.
+        """
+
+        self.pipeline = pipeline
+        self.model = self.pipeline.model
+        self.tokenizer = self.pipeline.tokenizer
+
+        self.model_name = self.pipeline.model.name_or_path
+
+    def _apply_chat_template(self, messages: list[dict[str, Any]]) -> str:
+        """
+        Applies a chat template to a list of messages.
+        If the tokenizer has a chat template, it uses that.
+        Otherwise, it falls back to a simple concatenation.
+
+        Args:
+            messages (list[dict]): A list of message dictionaries, e.g.,
+                                   [{"role": "user", "content": "Hello!"}, ...]
+
+        Returns:
+            The formatted prompt string ready for model input.
+        """
+
+        if hasattr(self.tokenizer, "apply_chat_template") and self.tokenizer.chat_template:
+            # Use the tokenizer's built-in chat template if available
+            # `tokenize=False` means it returns a string, not token IDs
+            return self.tokenizer.apply_chat_template(  # type: ignore[no-any-return]
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+        else:
+            # Fallback to a simple concatenation for models without a specific chat template
+            # This is a basic example; real chat models often need specific formatting.
+            prompt = ""
+            for message in messages:
+                role = message.get("role", "user")
+                content = message.get("content", "")
+                if role == "system":
+                    prompt += f"System: {content}\n"
+                elif role == "user":
+                    prompt += f"User: {content}\n"
+                elif role == "assistant":
+                    prompt += f"Assistant: {content}\n"
+            prompt += "Assistant:"  # Indicate that the assistant should respond
+            return prompt
+
+    def generate_chat_completion(
+        self,
+        messages: list[dict[str, Any]],
+        max_completion_tokens: Optional[int] = None,
+        stream: Optional[bool] = False,
+        stop_strings: Optional[list[str]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        frequency_penalty: Optional[float] = None,
+        presence_penalty: Optional[float] = None,
+        n: int = 1,
+    ) -> dict[str, Any]:
+        """
+        Generates a chat completion response in an OpenAI-compatible format.
+
+        Args:
+            messages (list[dict]): A list of message dictionaries, e.g.,
+                                   [{"role": "system", "content": "You are a helpful assistant."},
+                                    {"role": "user", "content": "What is deep learning?"}]
+            max_completion_tokens (int): The maximum number of completion tokens to generate.
+            stop_strings (list[str]): A list of strings to stop generation.
+            temperature (float): The temperature for sampling.
+            top_p (float): The top-p value for sampling.
+            stream (bool): Whether to stream the generation.
+            frequency_penalty (float): The frequency penalty for sampling.
+            presence_penalty (float): The presence penalty for sampling.
+            n (int): The number of samples to generate.
+
+        Returns:
+            dict: An OpenAI-compatible dictionary representing the chat completion.
+        """
+        # Apply chat template to convert messages into a single prompt string
+
+        prompt_text = self._apply_chat_template(messages)
+
+        # Tokenize the prompt
+        inputs = self.tokenizer(
+            prompt_text,
+            return_tensors="pt",
+            padding=True,
+        )
+        prompt_tokens = inputs.input_ids.shape[1]
+
+        from transformers import GenerationConfig
+
+        generation_config = GenerationConfig(
+            max_new_tokens=max_completion_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            pad_token_id=self.tokenizer.pad_token_id,
+            eos_token_id=self.tokenizer.eos_token_id,
+            stop_strings=stop_strings,
+            stream=stream,
+            repetition_penalty=frequency_penalty,
+            diversity_penalty=presence_penalty if n > 1 else None,
+            num_return_sequences=n,
+            num_beams=max(2, n),  # must be >1
+            num_beam_groups=max(2, n) if presence_penalty else 1,
+        )
+
+        # Generate text
+        output_ids = self.model.generate(
+            inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            generation_config=generation_config,
+        )
+
+        generated_texts = []
+        completion_tokens = 0
+        total_tokens = prompt_tokens
+        for output_id in output_ids:
+            # The output_ids include the input prompt
+            # Decode the generated text, excluding the input prompt
+            # so we slice to get only new tokens
+            generated_tokens = output_id[prompt_tokens:]
+            generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True)
+            generated_texts.append(generated_text)
+
+            # Calculate completion tokens
+            completion_tokens += len(generated_tokens)
+            total_tokens += len(generated_tokens)
+
+        choices = []
+        for i, generated_text in enumerate(generated_texts):
+            choices.append(
+                {
+                    "index": i,
+                    "message": {"role": "assistant", "content": generated_text},
+                    "logprobs": None,  # Not directly supported in this basic implementation
+                    "finish_reason": "stop",  # Assuming stop for simplicity
+                }
+            )
+
+        # Construct OpenAI-compatible response
+        response = {
+            "id": f"chatcmpl-{uuid.uuid4().hex}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": self.model_name,
+            "choices": choices,
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": total_tokens,
+            },
+        }
+        return response

snowflake/ml/model/_packager/model_handlers/sklearn.py

@@ -386,7 +386,9 @@ class SKLModelHandler(_base.BaseModelHandler[Union["sklearn.base.BaseEstimator",
             predictor = model[-1] if isinstance(model, sklearn.pipeline.Pipeline) else model
             try:
                 explainer = shap.Explainer(predictor, transformed_bg_data)
-                return handlers_utils.convert_explanations_to_2D_df(model, explainer(transformed_data).values)
+                return handlers_utils.convert_explanations_to_2D_df(model, explainer(transformed_data).values).astype(
+                    np.float64, errors="ignore"
+                )
             except TypeError:
                 if isinstance(data, pd.DataFrame):
                     dtype_map = {spec.name: spec.as_dtype(force_numpy_dtype=True) for spec in input_specs}

snowflake/ml/model/_packager/model_runtime/_snowml_inference_alternative_requirements.py

@@ -14,7 +14,7 @@ REQUIREMENTS = [
     "packaging>=20.9,<25",
     "pandas>=2.1.4,<3",
     "platformdirs<5",
-    "pyarrow",
+    "pyarrow<19.0.0",
     "pydantic>=2.8.2, <3",
     "pyjwt>=2.0.0, <3",
     "pytimeparse>=1.1.8,<2",
@@ -22,10 +22,10 @@ REQUIREMENTS = [
     "requests",
     "retrying>=1.3.3,<2",
     "s3fs>=2024.6.1,<2026",
-    "scikit-learn<1.6",
+    "scikit-learn<1.7",
     "scipy>=1.9,<2",
     "shap>=0.46.0,<1",
-    "snowflake-connector-python>=3.15.0,<4",
+    "snowflake-connector-python>=3.16.0,<4",
     "snowflake-snowpark-python>=1.17.0,<2,!=1.26.0",
     "snowflake.core>=1.0.2,<2",
     "sqlparse>=0.4,<1",

snowflake/ml/model/_signatures/snowpark_handler.py

@@ -84,7 +84,7 @@ class SnowparkDataFrameHandler(base_handler.BaseDataHandler[snowflake.snowpark.D
             return json.loads(x)
 
         for field in data.schema.fields:
-            if isinstance(field.datatype, spt.ArrayType):
+            if isinstance(field.datatype, (spt.ArrayType, spt.MapType, spt.StructType)):
                 df_local[identifier.get_unescaped_names(field.name)] = df_local[
                     identifier.get_unescaped_names(field.name)
                 ].map(load_if_not_null)

snowflake/ml/model/_signatures/utils.py

@@ -104,7 +104,10 @@ def rename_pandas_df(data: pd.DataFrame, features: Sequence[core.BaseFeatureSpec
     return data
 
 
-def huggingface_pipeline_signature_auto_infer(task: str, params: dict[str, Any]) -> Optional[core.ModelSignature]:
+def huggingface_pipeline_signature_auto_infer(
+    task: str,
+    params: dict[str, Any],
+) -> Optional[core.ModelSignature]:
     # Text
 
     # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.ConversationalPipeline
@@ -297,7 +300,6 @@ def huggingface_pipeline_signature_auto_infer(task: str, params: dict[str, Any])
                 )
             ],
         )
-
     # https://huggingface.co/docs/transformers/en/main_classes/pipelines#transformers.Text2TextGenerationPipeline
     if task == "text2text-generation":
         if params.get("return_tensors", False):

snowflake/ml/model/inference_engine.py

@@ -0,0 +1,5 @@
+import enum
+
+
+class InferenceEngine(enum.Enum):
+    VLLM = "vllm"

snowflake/ml/model/models/huggingface_pipeline.py

@@ -258,7 +258,7 @@ class HuggingFacePipelineModel:
         # model_version_impl.create_service parameters
         service_name: str,
         service_compute_pool: str,
-        image_repo: str,
+        image_repo: Optional[str] = None,
         image_build_compute_pool: Optional[str] = None,
         ingress_enabled: bool = False,
         max_instances: int = 1,
@@ -282,7 +282,8 @@ class HuggingFacePipelineModel:
             comment: Comment for the model. Defaults to None.
             service_name: The name of the service to create.
             service_compute_pool: The compute pool for the service.
-            image_repo: The name of the image repository.
+            image_repo: The name of the image repository. This can be None, in that case a default hidden image
+                repository will be used.
             image_build_compute_pool: The name of the compute pool used to build the model inference image. It uses
             the service compute pool if None.
             ingress_enabled: Whether ingress is enabled. Defaults to False.
@@ -356,7 +357,7 @@ class HuggingFacePipelineModel:
                         else sql_identifier.SqlIdentifier(service_compute_pool)
                     ),
                     service_compute_pool_name=sql_identifier.SqlIdentifier(service_compute_pool),
-                    image_repo=image_repo,
+                    image_repo_name=image_repo,
                     ingress_enabled=ingress_enabled,
                     max_instances=max_instances,
                     cpu_requests=cpu_requests,

snowflake/ml/model/openai_signatures.py

@@ -0,0 +1,57 @@
+from snowflake.ml.model._signatures import core
+
+_OPENAI_CHAT_SIGNATURE_SPEC = core.ModelSignature(
+    inputs=[
+        core.FeatureGroupSpec(
+            name="messages",
+            specs=[
+                core.FeatureSpec(name="content", dtype=core.DataType.STRING),
+                core.FeatureSpec(name="name", dtype=core.DataType.STRING),
+                core.FeatureSpec(name="role", dtype=core.DataType.STRING),
+                core.FeatureSpec(name="title", dtype=core.DataType.STRING),
+            ],
+            shape=(-1,),
+        ),
+        core.FeatureSpec(name="temperature", dtype=core.DataType.DOUBLE),
+        core.FeatureSpec(name="max_completion_tokens", dtype=core.DataType.INT64),
+        core.FeatureSpec(name="stop", dtype=core.DataType.STRING, shape=(-1,)),
+        core.FeatureSpec(name="n", dtype=core.DataType.INT32),
+        core.FeatureSpec(name="stream", dtype=core.DataType.BOOL),
+        core.FeatureSpec(name="top_p", dtype=core.DataType.DOUBLE),
+        core.FeatureSpec(name="frequency_penalty", dtype=core.DataType.DOUBLE),
+        core.FeatureSpec(name="presence_penalty", dtype=core.DataType.DOUBLE),
+    ],
+    outputs=[
+        core.FeatureSpec(name="id", dtype=core.DataType.STRING),
+        core.FeatureSpec(name="object", dtype=core.DataType.STRING),
+        core.FeatureSpec(name="created", dtype=core.DataType.FLOAT),
+        core.FeatureSpec(name="model", dtype=core.DataType.STRING),
+        core.FeatureGroupSpec(
+            name="choices",
+            specs=[
+                core.FeatureSpec(name="index", dtype=core.DataType.INT32),
+                core.FeatureGroupSpec(
+                    name="message",
+                    specs=[
+                        core.FeatureSpec(name="content", dtype=core.DataType.STRING),
+                        core.FeatureSpec(name="name", dtype=core.DataType.STRING),
+                        core.FeatureSpec(name="role", dtype=core.DataType.STRING),
+                    ],
+                ),
+                core.FeatureSpec(name="logprobs", dtype=core.DataType.STRING),
+                core.FeatureSpec(name="finish_reason", dtype=core.DataType.STRING),
+            ],
+            shape=(-1,),
+        ),
+        core.FeatureGroupSpec(
+            name="usage",
+            specs=[
+                core.FeatureSpec(name="completion_tokens", dtype=core.DataType.INT32),
+                core.FeatureSpec(name="prompt_tokens", dtype=core.DataType.INT32),
+                core.FeatureSpec(name="total_tokens", dtype=core.DataType.INT32),
+            ],
+        ),
+    ],
+)
+
+OPENAI_CHAT_SIGNATURE = {"__call__": _OPENAI_CHAT_SIGNATURE_SPEC}