mlflow.types

The mlflow.types module defines data types and utilities to be used by other mlflow components to describe interface independent of other frameworks or languages.

class mlflow.types.ColSpec(type: Union[mlflow.types.schema.DataType, mlflow.types.schema.Array, mlflow.types.schema.Object, str], name: Optional[str] = None, optional: Optional[bool] = None, required: Optional[bool] = None)

Bases: object

Specification of name and type of a single column in a dataset.

classmethod from_json_dict(**kwargs)

Deserialize from a json loaded dictionary. The dictionary is expected to contain type and optional name and required keys.

property name

The column name or None if the columns is unnamed.

property optional

Whether this column is optional.

Warning

Deprecated. optional is deprecated in favor of required.

Note

Experimental: This property may change or be removed in a future release without warning.

property required

Whether this column is required.

Note

Experimental: This property may change or be removed in a future release without warning.

property type

The column data type.

class mlflow.types.DataType(value)

Bases: enum.Enum

MLflow data types.

binary = 7

Sequence of raw bytes.

boolean = 1

Logical data (True, False) .

datetime = 8

64b datetime data.

double = 5

64b floating point numbers.

float = 4

32b floating point numbers.

integer = 2

32b signed integer numbers.

long = 3

64b signed integer numbers.

string = 6

Text data.

to_numpy() → numpy.dtype

Get equivalent numpy data type.

to_pandas() → numpy.dtype

Get equivalent pandas data type.

to_python()

Get equivalent python data type.

class mlflow.types.ParamSchema(params: List[mlflow.types.schema.ParamSpec])

Bases: object

Note

Experimental: This class may change or be removed in a future release without warning.

Specification of parameters applicable to the model. ParamSchema is represented as a list of ParamSpec.

classmethod from_json(json_str: str)

Deserialize from a json string.

property params

Representation of ParamSchema as a list of ParamSpec.

to_dict() → List[Dict[str, Any]]

Serialize into a jsonable dictionary.

to_json() → str

Serialize into json string.

class mlflow.types.ParamSpec(name: str, dtype: Union[mlflow.types.schema.DataType, str], default: Optional[Union[mlflow.types.schema.DataType, List[mlflow.types.schema.DataType]]], shape: Optional[Tuple[int, …]] = None)

Bases: object

Note

Experimental: This class may change or be removed in a future release without warning.

Specification used to represent parameters for the model.

class ParamSpecTypedDict(*args, **kwargs)

Bases: dict

property default

Default value of the parameter.

property dtype

The parameter data type.

classmethod enforce_param_datatype(name, value, dtype: mlflow.types.schema.DataType)

Enforce the value matches the data type.

The following type conversions are allowed:

1. int -> long, float, double

2. long -> float, double

3. float -> double

4. any -> datetime (try conversion)

Any other type mismatch will raise error.

Parameters

• name – parameter name

• value – parameter value

• t – expected data type

classmethod from_json_dict(**kwargs)

Deserialize from a json loaded dictionary. The dictionary is expected to contain name, type and default keys.

property name

The name of the parameter.

property shape

The parameter shape. If shape is None, the parameter is a scalar.

classmethod validate_type_and_shape(spec: str, value: Optional[Union[mlflow.types.schema.DataType, List[mlflow.types.schema.DataType]]], value_type: mlflow.types.schema.DataType, shape: Optional[Tuple[int, …]])

Validate that the value has the expected type and shape.

class mlflow.types.Schema(inputs: List[Union[mlflow.types.schema.ColSpec, mlflow.types.schema.TensorSpec]])

Bases: object

Specification of a dataset.

Schema is represented as a list of ColSpec or TensorSpec. A combination of ColSpec and TensorSpec is not allowed.

The dataset represented by a schema can be named, with unique non empty names for every input. In the case of ColSpec, the dataset columns can be unnamed with implicit integer index defined by their list indices. Combination of named and unnamed data inputs are not allowed.

as_spark_schema()

Convert to Spark schema. If this schema is a single unnamed column, it is converted directly the corresponding spark data type, otherwise it’s returned as a struct (missing column names are filled with an integer sequence). Unsupported by TensorSpec.

classmethod from_json(json_str: str)

Deserialize from a json string.

has_input_names() → bool

Return true iff this schema declares names, false otherwise.

input_dict() → Dict[str, Union[mlflow.types.schema.ColSpec, mlflow.types.schema.TensorSpec]]

Maps column names to inputs, iff this schema declares names.

input_names() → List[Union[str, int]]

Get list of data names or range of indices if the schema has no names.

input_types() → List[Union[mlflow.types.schema.DataType, numpy.dtype, mlflow.types.schema.Array, mlflow.types.schema.Object]]

Get types for each column in the schema.

input_types_dict() → Dict[str, Union[mlflow.types.schema.DataType, numpy.dtype, mlflow.types.schema.Array, mlflow.types.schema.Object]]

Maps column names to types, iff this schema declares names.

property inputs

Representation of a dataset that defines this schema.

is_tensor_spec() → bool

Return true iff this schema is specified using TensorSpec

numpy_types() → List[numpy.dtype]

Convenience shortcut to get the datatypes as numpy types.

optional_input_names() → List[Union[str, int]]

Note

Experimental: This function may change or be removed in a future release without warning.

Get list of optional data names or range of indices if schema has no names.

pandas_types() → List[numpy.dtype]

Convenience shortcut to get the datatypes as pandas types. Unsupported by TensorSpec.

required_input_names() → List[Union[str, int]]

Get list of required data names or range of indices if schema has no names.

to_dict() → List[Dict[str, Any]]

Serialize into a jsonable dictionary.

to_json() → str

Serialize into json string.

class mlflow.types.TensorSpec(type: numpy.dtype, shape: Union[tuple, list], name: Optional[str] = None)

Bases: object

Specification used to represent a dataset stored as a Tensor.

classmethod from_json_dict(**kwargs)

Deserialize from a json loaded dictionary. The dictionary is expected to contain type and tensor-spec keys.

property name

The tensor name or None if the tensor is unnamed.

property required

Whether this tensor is required.

Note

Experimental: This property may change or be removed in a future release without warning.

property shape

The tensor shape

property type

A unique character code for each of the 21 different numpy built-in types. See https://numpy.org/devdocs/reference/generated/numpy.dtype.html#numpy.dtype for details.

class mlflow.types.llm.ChatChoice(index: int, message: mlflow.types.llm.ChatMessage, finish_reason: str)

A single chat response generated by the model.

Parameters

• index (int) – The index of the response in the list of responses.

• message (ChatMessage) – The message that was generated.

• finish_reason (str) – The reason why generation stopped.

class mlflow.types.llm.ChatMessage(role: str, content: str, name: Optional[str] = None)

A message in a chat request or response.

Parameters

• role (str) – The role of the entity that sent the message (e.g. "user", "system").

• content (str) – The content of the message.

• name (str) – The name of the entity that sent the message. Optional.

class mlflow.types.llm.ChatParams(temperature: float = 1.0, max_tokens: Optional[int] = None, stop: Optional[List[str]] = None, n: int = 1, stream: bool = False)

Common parameters used for chat inference

Parameters

• temperature (float) – A param used to control randomness and creativity during inference. Optional, defaults to 1.0

• max_tokens (int) – The maximum number of new tokens to generate. Optional, defaults to None (unlimited)

• stop (List[str]) – A list of tokens at which to stop generation. Optional, defaults to None

• n (int) – The number of responses to generate. Optional, defaults to 1

• stream (bool) – Whether to stream back responses as they are generated. Optional, defaults to False

class mlflow.types.llm.ChatRequest(temperature: float = 1.0, max_tokens: Optional[int] = None, stop: Optional[List[str]] = None, n: int = 1, stream: bool = False, messages: List[mlflow.types.llm.ChatMessage] = <factory>)

Format of the request object expected by the chat endpoint.

Parameters

• messages (List[ChatMessage]) – A list of ChatMessage that will be passed to the model. Optional, defaults to empty list ([])

• temperature (float) – A param used to control randomness and creativity during inference. Optional, defaults to 1.0

• max_tokens (int) – The maximum number of new tokens to generate. Optional, defaults to None (unlimited)

• stop (List[str]) – A list of tokens at which to stop generation. Optional, defaults to None

• n (int) – The number of responses to generate. Optional, defaults to 1

• stream (bool) – Whether to stream back responses as they are generated. Optional, defaults to False

class mlflow.types.llm.ChatResponse(id: str, model: str, choices: List[mlflow.types.llm.ChatChoice], usage: mlflow.types.llm.TokenUsageStats, object: Literal[chat.completion] = 'chat.completion', created: int = <factory>)

The full response object returned by the chat endpoint.

Parameters

• id (str) – The ID of the response.

• object (str) – The object type.

• created (int) – The time the response was created. Optional, defaults to the current time.

• model (str) – The name of the model used.

• choices (List[ChatChoice]) – A list of ChatChoice objects containing the generated responses

• usage (TokenUsageStats) – An object describing the tokens used by the request.

class mlflow.types.llm.TokenUsageStats(prompt_tokens: int, completion_tokens: int, total_tokens: int)

Stats about the number of tokens used during inference.

Parameters

• prompt_tokens (int) – The number of tokens in the prompt.

• completion_tokens (int) – The number of tokens in the generated completion.

• total_tokens (int) – The total number of tokens used.