wehub-resource-sync
55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
|
|
|
|
if TYPE_CHECKING:
|
|
import torch
|
|
|
|
_TORCH_DTYPE_KEY = "torch_dtype"
|
|
|
|
|
|
def _extract_torch_dtype_if_set(pipeline) -> torch.dtype | None:
|
|
"""
|
|
Extract the torch datatype argument if set and return as a string encoded value.
|
|
"""
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
# If torch is not installed, safe to assume the model doesn't have a custom torch_dtype
|
|
return None
|
|
|
|
# Check model dtype as pipeline's torch_dtype field doesn't always reflect the model's dtype
|
|
model_dtype = pipeline.model.dtype if hasattr(pipeline.model, "dtype") else None
|
|
|
|
# If the underlying model is PyTorch model, dtype must be a torch.dtype instance
|
|
return model_dtype if isinstance(model_dtype, torch.dtype) else None
|
|
|
|
|
|
def _deserialize_torch_dtype(dtype_str: str) -> torch.dtype:
|
|
"""
|
|
Convert the string-encoded `torch_dtype` pipeline argument back to the correct `torch.dtype`
|
|
instance value for applying to a loaded pipeline instance.
|
|
"""
|
|
try:
|
|
import torch
|
|
except ImportError as e:
|
|
raise MlflowException(
|
|
"Unable to determine if the value supplied by the argument "
|
|
"torch_dtype is valid since torch is not installed.",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
) from e
|
|
|
|
dtype_str = dtype_str.removeprefix("torch.")
|
|
|
|
dtype = getattr(torch, dtype_str, None)
|
|
if isinstance(dtype, torch.dtype):
|
|
return dtype
|
|
|
|
raise MlflowException(
|
|
f"The value '{dtype_str}' is not a valid torch.dtype",
|
|
error_code=INVALID_PARAMETER_VALUE,
|
|
)
|