snowflake-ml-python 1.11.0__py3-none-any.whl → 1.13.0__py3-none-any.whl

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

@@ -1,6 +1,9 @@
 import json
 import logging
 import os
+import shutil
+import time
+import uuid
 import warnings
 from typing import TYPE_CHECKING, Any, Callable, Optional, Union, cast, final
 
@@ -11,7 +14,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
@@ -81,6 +89,7 @@ class HuggingFacePipelineHandler(
     _HANDLER_MIGRATOR_PLANS: dict[str, type[base_migrator.BaseModelHandlerMigrator]] = {}
 
     MODEL_BLOB_FILE_OR_DIR = "model"
+    MODEL_PICKLE_FILE = "snowml_huggingface_pipeline.pkl"
     ADDITIONAL_CONFIG_FILE = "pipeline_config.pt"
     DEFAULT_TARGET_METHODS = ["__call__"]
     IS_AUTO_SIGNATURE = True
@@ -151,7 +160,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(
@@ -189,6 +201,7 @@ class HuggingFacePipelineHandler(
         model_blob_path = os.path.join(model_blobs_dir_path, name)
         os.makedirs(model_blob_path, exist_ok=True)
 
+        is_repo_downloaded = False
         if type_utils.LazyType("transformers.Pipeline").isinstance(model):
             save_path = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
             model.save_pretrained(  # type:ignore[attr-defined]
@@ -214,11 +227,22 @@ class HuggingFacePipelineHandler(
             ) as f:
                 cloudpickle.dump(pipeline_params, f)
         else:
+            model_blob_file_or_dir = os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR)
+            model_blob_pickle_file = os.path.join(model_blob_file_or_dir, cls.MODEL_PICKLE_FILE)
+            os.makedirs(model_blob_file_or_dir, exist_ok=True)
             with open(
-                os.path.join(model_blob_path, cls.MODEL_BLOB_FILE_OR_DIR),
+                model_blob_pickle_file,
                 "wb",
             ) as f:
                 cloudpickle.dump(model, f)
+            if model.repo_snapshot_dir:
+                logger.info("model's repo_snapshot_dir is available, copying snapshot")
+                shutil.copytree(
+                    model.repo_snapshot_dir,
+                    model_blob_file_or_dir,
+                    dirs_exist_ok=True,
+                )
+                is_repo_downloaded = True
 
         base_meta = model_blob_meta.ModelBlobMeta(
             name=name,
@@ -226,13 +250,12 @@ class HuggingFacePipelineHandler(
             handler_version=cls.HANDLER_VERSION,
             path=cls.MODEL_BLOB_FILE_OR_DIR,
             options=model_meta_schema.HuggingFacePipelineModelBlobOptions(
-                {
-                    "task": task,
-                    "batch_size": batch_size if batch_size is not None else 1,
-                    "has_tokenizer": has_tokenizer,
-                    "has_feature_extractor": has_feature_extractor,
-                    "has_image_preprocessor": has_image_preprocessor,
-                }
+                task=task,
+                batch_size=batch_size if batch_size is not None else 1,
+                has_tokenizer=has_tokenizer,
+                has_feature_extractor=has_feature_extractor,
+                has_image_preprocessor=has_image_preprocessor,
+                is_repo_downloaded=is_repo_downloaded,
             ),
         )
         model_meta.models[name] = base_meta
@@ -276,6 +299,27 @@ class HuggingFacePipelineHandler(
 
         return device_config
 
+    @staticmethod
+    def _load_pickle_model(
+        pickle_file: str,
+        **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+    ) -> huggingface_pipeline.HuggingFacePipelineModel:
+        with open(pickle_file, "rb") as f:
+            m = cloudpickle.load(f)
+        assert isinstance(m, huggingface_pipeline.HuggingFacePipelineModel)
+        torch_dtype: Optional[str] = None
+        device_config = None
+        if getattr(m, "device", None) is None and getattr(m, "device_map", None) is None:
+            device_config = HuggingFacePipelineHandler._get_device_config(**kwargs)
+            m.__dict__.update(device_config)
+
+        if getattr(m, "torch_dtype", None) is None and kwargs.get("use_gpu", False):
+            torch_dtype = "auto"
+            m.__dict__.update(torch_dtype=torch_dtype)
+        else:
+            m.__dict__.update(torch_dtype=None)
+        return m
+
     @classmethod
     def load_model(
         cls,
@@ -300,7 +344,13 @@ class HuggingFacePipelineHandler(
             raise ValueError("Missing field `batch_size` in model blob metadata for type `huggingface_pipeline`")
 
         model_blob_file_or_dir_path = os.path.join(model_blob_path, model_blob_filename)
-        if os.path.isdir(model_blob_file_or_dir_path):
+        is_repo_downloaded = model_blob_options.get("is_repo_downloaded", False)
+
+        def _create_pipeline_from_dir(
+            model_blob_file_or_dir_path: str,
+            model_blob_options: model_meta_schema.HuggingFacePipelineModelBlobOptions,
+            **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+        ) -> "transformers.Pipeline":
             import transformers
 
             additional_pipeline_params = {}
@@ -320,7 +370,7 @@ class HuggingFacePipelineHandler(
             ) as f:
                 pipeline_params = cloudpickle.load(f)
 
-            device_config = cls._get_device_config(**kwargs)
+            device_config = HuggingFacePipelineHandler._get_device_config(**kwargs)
 
             m = transformers.pipeline(
                 model_blob_options["task"],
@@ -349,18 +399,59 @@ class HuggingFacePipelineHandler(
                 m.tokenizer.chat_template = DEFAULT_CHAT_TEMPLATE
 
             m.__dict__.update(pipeline_params)
+            return m
 
+        def _create_pipeline_from_model(
+            model_blob_file_or_dir_path: str,
+            m: huggingface_pipeline.HuggingFacePipelineModel,
+            **kwargs: Unpack[model_types.HuggingFaceLoadOptions],
+        ) -> "transformers.Pipeline":
+            import transformers
+
+            return transformers.pipeline(
+                m.task,
+                model=model_blob_file_or_dir_path,
+                trust_remote_code=m.trust_remote_code,
+                torch_dtype=getattr(m, "torch_dtype", None),
+                revision=m.revision,
+                # pass device or device_map when creating the pipeline
+                **HuggingFacePipelineHandler._get_device_config(**kwargs),
+                # pass other model_kwargs to transformers.pipeline.from_pretrained method
+                **m.model_kwargs,
+            )
+
+        if os.path.isdir(model_blob_file_or_dir_path) and not is_repo_downloaded:
+            # the logged model is a transformers.Pipeline object
+            # weights of the model are saved in the directory
+            return _create_pipeline_from_dir(model_blob_file_or_dir_path, model_blob_options, **kwargs)
         else:
-            assert os.path.isfile(model_blob_file_or_dir_path)
-            with open(model_blob_file_or_dir_path, "rb") as f:
-                m = cloudpickle.load(f)
-            assert isinstance(m, huggingface_pipeline.HuggingFacePipelineModel)
-            if getattr(m, "device", None) is None and getattr(m, "device_map", None) is None:
-                m.__dict__.update(cls._get_device_config(**kwargs))
-
-            if getattr(m, "torch_dtype", None) is None and kwargs.get("use_gpu", False):
-                m.__dict__.update(torch_dtype="auto")
-        return m
+            # case 1: LEGACY logging, repo snapshot is not logged
+            if os.path.isfile(model_blob_file_or_dir_path):
+                # LEGACY logging that had model as a pickle file in the model blob directory
+                # the logged model is a huggingface_pipeline.HuggingFacePipelineModel object
+                # the model_blob_file_or_dir_path is the pickle file that holds
+                # the huggingface_pipeline.HuggingFacePipelineModel object
+                # the snapshot of the repo is not logged
+                return cls._load_pickle_model(model_blob_file_or_dir_path)
+            else:
+                assert os.path.isdir(model_blob_file_or_dir_path)
+                # the logged model is a huggingface_pipeline.HuggingFacePipelineModel object
+                # the pickle_file holds the huggingface_pipeline.HuggingFacePipelineModel object
+                pickle_file = os.path.join(model_blob_file_or_dir_path, cls.MODEL_PICKLE_FILE)
+                m = cls._load_pickle_model(pickle_file)
+
+                # case 2: logging without the snapshot of the repo
+                if not is_repo_downloaded:
+                    # we return the huggingface_pipeline.HuggingFacePipelineModel object
+                    return m
+                # case 3: logging with the snapshot of the repo
+                else:
+                    # the model_blob_file_or_dir_path is the directory that holds
+                    # weights of the model from `huggingface_hub.snapshot_download`
+                    # the huggingface_pipeline.HuggingFacePipelineModel object is logged
+                    # with a snapshot of the repo, we create a transformers.Pipeline object
+                    # by reading the snapshot directory
+                    return _create_pipeline_from_model(model_blob_file_or_dir_path, m, **kwargs)
 
     @classmethod
     def convert_as_custom_model(
@@ -401,6 +492,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 +646,171 @@ 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,
+        ).to(self.model.device)
+        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,
+            do_sample=False,
+        )
+
+        # 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_handlers/xgboost.py

@@ -229,6 +229,11 @@ class XGBModelHandler(_base.BaseModelHandler[Union["xgboost.Booster", "xgboost.X
                     enable_categorical = False
                     for col, d_type in X.dtypes.items():
                         if pd.api.extensions.ExtensionDtype.is_dtype(d_type):
+                            if pd.CategoricalDtype.is_dtype(d_type):
+                                enable_categorical = True
+                            elif isinstance(d_type, pd.StringDtype):
+                                X[col] = X[col].astype("category")
+                                enable_categorical = True
                             continue
                         if not np.issubdtype(d_type, np.number):
                             # categorical columns are converted to numpy's str dtype

snowflake/ml/model/_packager/model_meta/model_meta_schema.py

@@ -51,6 +51,7 @@ class HuggingFacePipelineModelBlobOptions(BaseModelBlobOptions):
     has_tokenizer: NotRequired[bool]
     has_feature_extractor: NotRequired[bool]
     has_image_preprocessor: NotRequired[bool]
+    is_repo_downloaded: NotRequired[Optional[bool]]
 
 
 class LightGBMModelBlobOptions(BaseModelBlobOptions):

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/models/huggingface_pipeline.py

@@ -28,6 +28,10 @@ class HuggingFacePipelineModel:
         token: Optional[str] = None,
         trust_remote_code: Optional[bool] = None,
         model_kwargs: Optional[dict[str, Any]] = None,
+        download_snapshot: bool = True,
+        # repo snapshot download args
+        allow_patterns: Optional[Union[list[str], str]] = None,
+        ignore_patterns: Optional[Union[list[str], str]] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -52,6 +56,9 @@ class HuggingFacePipelineModel:
                 Defaults to None.
             model_kwargs: Additional dictionary of keyword arguments passed along to the model's `from_pretrained(...,`.
                 Defaults to None.
+            download_snapshot: Whether to download the HuggingFace repository. Defaults to True.
+            allow_patterns: If provided, only files matching at least one pattern are downloaded.
+            ignore_patterns: If provided, files matching any of the patterns are not downloaded.
             kwargs: Additional keyword arguments passed along to the specific pipeline init (see the documentation for
                 the corresponding pipeline class for possible values).
 
@@ -220,6 +227,21 @@ class HuggingFacePipelineModel:
                     stacklevel=2,
                 )
 
+        repo_snapshot_dir: Optional[str] = None
+        if download_snapshot:
+            try:
+                from huggingface_hub import snapshot_download
+
+                repo_snapshot_dir = snapshot_download(
+                    repo_id=model,
+                    revision=revision,
+                    token=token,
+                    allow_patterns=allow_patterns,
+                    ignore_patterns=ignore_patterns,
+                )
+            except ImportError:
+                logger.info("huggingface_hub package is not installed, skipping snapshot download")
+
         # ==== End pipeline logic from transformers ====
 
         self.task = normalized_task
@@ -229,6 +251,7 @@ class HuggingFacePipelineModel:
         self.trust_remote_code = trust_remote_code
         self.model_kwargs = model_kwargs
         self.tokenizer = tokenizer
+        self.repo_snapshot_dir = repo_snapshot_dir
         self.__dict__.update(kwargs)
 
     @telemetry.send_api_usage_telemetry(

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}