6d8c18cd1a* Replace standalone Studio wording with Unsloth Replace the single word Studio with Unsloth wherever it is used as shorthand for Unsloth Studio in docs, CLI output, UI strings, i18n locales, workflow display names, comments and docstrings. Kept unchanged: the full name Unsloth Studio, third party product names (LM Studio, Visual Studio, Mac Studio), feature names (Recipe Studio, Fine-tuning Studio and its translations), and all identifiers such as env vars, commands, paths and filenames. * Address review feedback on the Studio wording rename Use "an" before Unsloth where the rename left the article as "a". Restore the split brand where Unsloth and Studio render as two halves of the full product name: the onboarding sidebar subtitle and the IPv6 localhost warning. Scope two messages to the full name Unsloth Studio where plain Unsloth was misleading: the AMD README bullet and the CLI studio setup error.
336 lines
12 KiB
Python
336 lines
12 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from utils.paths import recipe_datasets_root
|
|
|
|
from .jsonable import to_jsonable
|
|
from .local_callable_validators import (
|
|
register_oxc_local_callable_validators,
|
|
split_oxc_local_callable_validators,
|
|
)
|
|
|
|
_IMAGE_CONTEXT_PATCHED = False
|
|
|
|
|
|
def _encode_bytes_to_base64(value: bytes | bytearray) -> str:
|
|
return base64.b64encode(bytes(value)).decode("utf-8")
|
|
|
|
|
|
def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None:
|
|
try:
|
|
path = Path(path_value)
|
|
candidates: list[Path] = []
|
|
if path.is_absolute():
|
|
candidates.append(path)
|
|
else:
|
|
if base_path:
|
|
candidates.append(Path(base_path) / path)
|
|
candidates.append(Path.cwd() / path)
|
|
|
|
for candidate in candidates:
|
|
if not candidate.exists() or not candidate.is_file():
|
|
continue
|
|
with candidate.open("rb") as f:
|
|
return _encode_bytes_to_base64(f.read())
|
|
except (OSError, TypeError, ValueError):
|
|
return None
|
|
return None
|
|
|
|
|
|
def _pil_image_to_base64(value: Any) -> str | None:
|
|
try:
|
|
from PIL.Image import Image as PILImage # type: ignore
|
|
except ImportError:
|
|
return None
|
|
if not isinstance(value, PILImage):
|
|
return None
|
|
buffer = io.BytesIO()
|
|
image_format = str(getattr(value, "format", "") or "").upper()
|
|
if image_format not in {"PNG", "JPEG", "JPG", "WEBP", "GIF"}:
|
|
image_format = "PNG"
|
|
value.save(buffer, format = image_format)
|
|
return _encode_bytes_to_base64(buffer.getvalue())
|
|
|
|
|
|
def _normalize_image_context_value(value: Any, *, base_path: str | None = None) -> Any:
|
|
if isinstance(value, str):
|
|
return value
|
|
|
|
if isinstance(value, (bytes, bytearray)):
|
|
return _encode_bytes_to_base64(value)
|
|
|
|
pil_base64 = _pil_image_to_base64(value)
|
|
if pil_base64 is not None:
|
|
return pil_base64
|
|
|
|
if isinstance(value, dict):
|
|
url = value.get("url")
|
|
if isinstance(url, str):
|
|
return url
|
|
|
|
image_url = value.get("image_url")
|
|
if isinstance(image_url, str):
|
|
return image_url
|
|
if isinstance(image_url, dict):
|
|
nested_url = image_url.get("url")
|
|
if isinstance(nested_url, str):
|
|
return nested_url
|
|
|
|
inline_data = value.get("data")
|
|
if isinstance(inline_data, str):
|
|
return inline_data
|
|
|
|
raw_bytes = value.get("bytes")
|
|
if isinstance(raw_bytes, (bytes, bytearray)):
|
|
return _encode_bytes_to_base64(raw_bytes)
|
|
if isinstance(raw_bytes, str) and raw_bytes.strip():
|
|
return raw_bytes
|
|
|
|
path_value = value.get("path")
|
|
if isinstance(path_value, str) and path_value.strip():
|
|
if as_base64 := _load_image_file_to_base64(path_value, base_path = base_path):
|
|
return as_base64
|
|
return path_value
|
|
|
|
return value
|
|
|
|
|
|
def _apply_data_designer_image_context_patch() -> None:
|
|
global _IMAGE_CONTEXT_PATCHED
|
|
if _IMAGE_CONTEXT_PATCHED:
|
|
return
|
|
|
|
try:
|
|
from data_designer.config.models import ImageContext # pyright: ignore[reportMissingImports]
|
|
except ImportError:
|
|
return
|
|
|
|
if getattr(ImageContext, "_unsloth_image_context_patch_applied", False):
|
|
_IMAGE_CONTEXT_PATCHED = True
|
|
return
|
|
|
|
original_auto_resolve = ImageContext._auto_resolve_context_value
|
|
|
|
def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any:
|
|
normalized = _normalize_image_context_value(context_value, base_path = base_path)
|
|
return original_auto_resolve(self, normalized, base_path)
|
|
|
|
ImageContext._auto_resolve_context_value = _patched_auto_resolve
|
|
setattr(ImageContext, "_unsloth_image_context_patch_applied", True)
|
|
_IMAGE_CONTEXT_PATCHED = True
|
|
|
|
|
|
def build_model_providers(recipe: dict[str, Any]):
|
|
from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports]
|
|
|
|
providers: list[ModelProvider] = []
|
|
for provider in recipe.get("model_providers", []):
|
|
api_key = provider.get("api_key")
|
|
api_key_env = provider.get("api_key_env")
|
|
if not api_key and api_key_env:
|
|
api_key = os.getenv(api_key_env)
|
|
providers.append(
|
|
ModelProvider(
|
|
name = provider["name"],
|
|
endpoint = provider["endpoint"],
|
|
provider_type = provider.get("provider_type", "openai"),
|
|
api_key = api_key,
|
|
extra_headers = provider.get("extra_headers"),
|
|
extra_body = provider.get("extra_body"),
|
|
)
|
|
)
|
|
|
|
return providers
|
|
|
|
|
|
def _recipe_has_llm_columns(recipe: dict[str, Any]) -> bool:
|
|
for column in recipe.get("columns", []):
|
|
if not isinstance(column, dict):
|
|
continue
|
|
column_type = column.get("column_type")
|
|
if isinstance(column_type, str) and column_type.startswith("llm-"):
|
|
return True
|
|
return False
|
|
|
|
|
|
def _validate_recipe_runtime_support(recipe: dict[str, Any], model_providers: list[Any]) -> None:
|
|
if _recipe_has_llm_columns(recipe) and not model_providers:
|
|
raise ValueError("Add a Provider connection block before running this recipe.")
|
|
|
|
|
|
def build_mcp_providers(recipe: dict[str, Any]) -> list:
|
|
from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports]
|
|
|
|
# Same gate as the chat MCP path: stdio providers spawn a local subprocess,
|
|
# so build them only when this host allows it (desktop / explicit opt-in).
|
|
from core.inference.mcp_client import stdio_mcp_enabled
|
|
|
|
stdio_allowed = stdio_mcp_enabled()
|
|
|
|
providers: list[MCPProvider | LocalStdioMCPProvider] = []
|
|
for provider in recipe.get("mcp_providers", []):
|
|
if not isinstance(provider, dict):
|
|
continue
|
|
provider_type = provider.get("provider_type")
|
|
if provider_type == "stdio":
|
|
if not stdio_allowed:
|
|
continue
|
|
env = provider.get("env")
|
|
if not isinstance(env, dict):
|
|
env = {}
|
|
args = provider.get("args")
|
|
if not isinstance(args, list):
|
|
args = []
|
|
providers.append(
|
|
LocalStdioMCPProvider(
|
|
name = str(provider.get("name", "")),
|
|
command = str(provider.get("command", "")),
|
|
args = [str(value) for value in args],
|
|
env = {str(key): str(value) for key, value in env.items()},
|
|
)
|
|
)
|
|
continue
|
|
|
|
if provider_type in {"sse", "streamable_http"}:
|
|
api_key = provider.get("api_key")
|
|
api_key_env = provider.get("api_key_env")
|
|
if not api_key and api_key_env:
|
|
api_key = os.getenv(str(api_key_env))
|
|
providers.append(
|
|
MCPProvider(
|
|
name = str(provider.get("name", "")),
|
|
endpoint = str(provider.get("endpoint", "")),
|
|
provider_type = str(provider_type),
|
|
api_key = str(api_key) if api_key else None,
|
|
)
|
|
)
|
|
return providers
|
|
|
|
|
|
def _strip_frontend_model_config_metadata(recipe: dict[str, Any]) -> dict[str, Any]:
|
|
model_configs = recipe.get("model_configs")
|
|
if not isinstance(model_configs, list):
|
|
return recipe
|
|
|
|
changed = False
|
|
next_model_configs: list[Any] = []
|
|
for model_config in model_configs:
|
|
if isinstance(model_config, dict) and "gguf_variant" in model_config:
|
|
next_model_config = dict(model_config)
|
|
next_model_config.pop("gguf_variant", None)
|
|
next_model_configs.append(next_model_config)
|
|
changed = True
|
|
continue
|
|
next_model_configs.append(model_config)
|
|
|
|
if not changed:
|
|
return recipe
|
|
|
|
return {
|
|
**recipe,
|
|
"model_configs": next_model_configs,
|
|
}
|
|
|
|
|
|
def build_config_builder(recipe: dict[str, Any]):
|
|
_apply_data_designer_image_context_patch()
|
|
from data_designer.config import DataDesignerConfigBuilder # pyright: ignore[reportMissingImports]
|
|
from data_designer.config.processors import ProcessorType # pyright: ignore[reportMissingImports]
|
|
|
|
recipe_core = {
|
|
key: value
|
|
for key, value in recipe.items()
|
|
if key not in {"model_providers", "mcp_providers"}
|
|
}
|
|
recipe_core = _strip_frontend_model_config_metadata(recipe_core)
|
|
recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators(recipe_core)
|
|
builder = DataDesignerConfigBuilder.from_config({"data_designer": recipe_core})
|
|
register_oxc_local_callable_validators(
|
|
builder = builder,
|
|
specs = oxc_local_callable_specs,
|
|
)
|
|
|
|
# DataDesignerConfigBuilder.from_config currently skips processors.
|
|
# Re-attach so drop_columns/schema_transform survive the API payload.
|
|
for processor in recipe_core.get("processors") or []:
|
|
if not isinstance(processor, dict):
|
|
continue
|
|
processor_type_raw = processor.get("processor_type")
|
|
if not isinstance(processor_type_raw, str):
|
|
continue
|
|
kwargs = {k: v for k, v in processor.items() if k != "processor_type"}
|
|
builder.add_processor(
|
|
processor_type = ProcessorType(processor_type_raw),
|
|
**kwargs,
|
|
)
|
|
|
|
return builder
|
|
|
|
|
|
def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = None):
|
|
_apply_data_designer_image_context_patch()
|
|
from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports]
|
|
|
|
if artifact_path is None:
|
|
# DataDesigner defaults to cwd/artifacts; packaged Unsloth can run with
|
|
# cwd=/, so keep default callers on Unsloth's writable recipe artifact root.
|
|
artifact_path = str(recipe_datasets_root())
|
|
|
|
recipe = _strip_frontend_model_config_metadata(recipe)
|
|
model_providers = build_model_providers(recipe)
|
|
_validate_recipe_runtime_support(recipe, model_providers)
|
|
|
|
# DataDesigner requires >=1 model provider even with no LLM columns; stub
|
|
# one so sampler/expression-only recipes run without a real provider.
|
|
if not model_providers:
|
|
from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports]
|
|
model_providers = [
|
|
ModelProvider(
|
|
name = "_unused",
|
|
endpoint = "http://localhost",
|
|
provider_type = "openai",
|
|
api_key = None,
|
|
)
|
|
]
|
|
|
|
return DataDesigner(
|
|
artifact_path = artifact_path,
|
|
model_providers = model_providers,
|
|
mcp_providers = build_mcp_providers(recipe),
|
|
)
|
|
|
|
|
|
def validate_recipe(recipe: dict[str, Any]) -> None:
|
|
builder = build_config_builder(recipe)
|
|
designer = create_data_designer(recipe)
|
|
designer.validate(builder)
|
|
|
|
|
|
def preview_recipe(
|
|
recipe: dict[str, Any], num_records: int
|
|
) -> tuple[list[dict[str, Any]], dict[str, Any] | None, dict[str, Any] | None]:
|
|
builder = build_config_builder(recipe)
|
|
designer = create_data_designer(recipe)
|
|
results = designer.preview(builder, num_records = num_records)
|
|
|
|
dataset: list[dict[str, Any]] = []
|
|
if results.dataset is not None:
|
|
raw_rows = results.dataset.to_dict(orient = "records")
|
|
dataset = [to_jsonable(row) for row in raw_rows]
|
|
|
|
artifacts = (
|
|
None if results.processor_artifacts is None else to_jsonable(results.processor_artifacts)
|
|
)
|
|
analysis = (
|
|
None if results.analysis is None else to_jsonable(results.analysis.model_dump(mode = "json"))
|
|
)
|
|
|
|
return dataset, artifacts, analysis
|