ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Final cleanup

Browse source
This commit is contained in:
Roland Tannous 2026-03-12 18:28:04 +00:00
commit 985d2e43ee
123 changed files with 7474 additions and 5805 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/backend/utils/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

7
studio/backend/utils/cache_cleanup.py
View file

@ -8,6 +8,7 @@ The unsloth_compiled_cache is created by unsloth_zoo/compiler.py during
FastModel.from_pretrained() and contains model-type-specific compiled Python
files. It should be cleared between model loads to avoid stale artefacts.
"""
import shutil
import structlog
from loggers import get_logger
@ -16,8 +17,8 @@ from pathlib import Path
logger = get_logger(__name__)
# Possible locations where unsloth_compiled_cache may appear
_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend
_PROJECT_ROOT = _BACKEND_DIR.parent.parent # repo root
_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend
_PROJECT_ROOT = _BACKEND_DIR.parent.parent # repo root
_CACHE_DIRS = [
_BACKEND_DIR / "unsloth_compiled_cache",
@ -31,4 +32,4 @@ def clear_unsloth_compiled_cache() -> None:
for cache_dir in _CACHE_DIRS:
if cache_dir.exists():
logger.info(f"Removing unsloth compiled cache: {cache_dir}")
shutil.rmtree(cache_dir, ignore_errors=True)
shutil.rmtree(cache_dir, ignore_errors = True)

28
studio/backend/utils/datasets/chat_templates.py
View file

@ -61,7 +61,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
try:
tokenizer = get_chat_template(
tokenizer,
chat_template=matched_template,
chat_template = matched_template,
)
except Exception as e:
logger.info(f"⚠️ Failed to apply Unsloth template '{matched_template}': {e}")
@ -80,7 +80,7 @@ def get_tokenizer_chat_template(tokenizer, model_name):
try:
tokenizer = get_chat_template(
tokenizer,
chat_template="chatml",
chat_template = "chatml",
)
except Exception as e:
logger.info(f"⚠️ Failed to apply default ChatML template: {e}")
@ -119,14 +119,14 @@ def get_dataset_info_summary(dataset_info):
def apply_chat_template_to_dataset(
dataset_info,
tokenizer,
model_name=None,
custom_prompt_template=None,
add_eos_token=False,
remove_bos_prefix=False,
custom_format_mapping=None,
auto_detect_mapping=True,
batch_size=1000,
num_proc=None,
model_name = None,
custom_prompt_template = None,
add_eos_token = False,
remove_bos_prefix = False,
custom_format_mapping = None,
auto_detect_mapping = True,
batch_size = 1000,
num_proc = None,
):
"""
Applies chat template to dataset based on its format.
@ -233,7 +233,7 @@ def apply_chat_template_to_dataset(
return result
try:
dataset = dataset.map(_apply_custom_mapping, batched=True, batch_size=batch_size)
dataset = dataset.map(_apply_custom_mapping, batched = True, batch_size = batch_size)
# Update to use conversations format
final_format = "chatml_conversations"
chat_column = "conversations"
@ -256,7 +256,7 @@ def apply_chat_template_to_dataset(
if not (hasattr(tokenizer, 'chat_template') and tokenizer.chat_template):
try:
from unsloth.chat_templates import get_chat_template
tokenizer = get_chat_template(tokenizer, chat_template="alpaca")
tokenizer = get_chat_template(tokenizer, chat_template = "alpaca")
logger.info(f"📝 Set alpaca chat template on tokenizer for model saving")
except Exception as e:
logger.info(f"⚠️ Could not set alpaca template on tokenizer: {e}")
@ -333,8 +333,8 @@ def apply_chat_template_to_dataset(
try:
text = tokenizer.apply_chat_template(
convo,
tokenize=False,
add_generation_prompt=False
tokenize = False,
add_generation_prompt = False
)
if remove_bos_prefix:

52
studio/backend/utils/datasets/data_collators.py
View file

@ -12,11 +12,10 @@ import torch
from dataclasses import dataclass
from typing import Any, List, Optional, Union
from loggers import get_logger
logger = get_logger(__name__)
@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
"""
@ -26,16 +25,23 @@ class DataCollatorSpeechSeq2SeqWithPadding:
masks padding in labels with -100, and strips leading BOS token.
Mirrors the collator from the Whisper.ipynb notebook.
"""
processor: Any
def __call__(self, features: List[dict]) -> dict:
input_features = [{"input_features": feature["input_features"]} for feature in features]
batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
input_features = [
{"input_features": feature["input_features"]} for feature in features
]
batch = self.processor.feature_extractor.pad(
input_features, return_tensors = "pt"
)
label_features = [{"input_ids": feature["labels"]} for feature in features]
labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
labels_batch = self.processor.tokenizer.pad(label_features, return_tensors = "pt")
labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
labels = labels_batch["input_ids"].masked_fill(
labels_batch.attention_mask.ne(1), -100
)
if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
labels = labels[:, 1:]
@ -54,6 +60,7 @@ class DeepSeekOCRDataCollator:
- Text tokenization
- Proper label masking for instruction fine-tuning
"""
processor: Any # Qwen2VLProcessor or similar
max_length: int = 2048
ignore_index: int = -100
@ -86,7 +93,7 @@ class DeepSeekOCRDataCollator:
for item in content:
if isinstance(item, dict) and item.get("type") == "image":
img = item.get("image")
if img is not None and hasattr(img, 'size'): # PIL Image
if img is not None and hasattr(img, "size"): # PIL Image
all_images.append(img)
# Process with the VL processor
@ -94,19 +101,19 @@ class DeepSeekOCRDataCollator:
# Qwen2VL style processing
texts = [
self.processor.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=False
msgs, tokenize = False, add_generation_prompt = False
)
for msgs in all_messages
]
# Process with images
inputs = self.processor(
text=texts,
images=all_images if all_images else None,
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.max_length,
text = texts,
images = all_images if all_images else None,
return_tensors = "pt",
padding = True,
truncation = True,
max_length = self.max_length,
)
# Create labels (mask input, keep output)
@ -134,6 +141,7 @@ class VLMDataCollator:
- LLaVA
- Other VL models with compatible processors
"""
processor: Any
max_length: int = 2048
ignore_index: int = -100
@ -163,26 +171,26 @@ class VLMDataCollator:
# Apply chat template
texts = [
self.processor.apply_chat_template(
msgs, tokenize=False, add_generation_prompt=False
msgs, tokenize = False, add_generation_prompt = False
)
for msgs in all_messages
]
# Process inputs
inputs = self.processor(
text=texts,
images=all_images if all_images else None,
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.max_length,
text = texts,
images = all_images if all_images else None,
return_tensors = "pt",
padding = True,
truncation = True,
max_length = self.max_length,
)
# Create labels
labels = inputs["input_ids"].clone()
# Mask padding
if hasattr(self.processor, 'tokenizer'):
if hasattr(self.processor, "tokenizer"):
pad_token_id = self.processor.tokenizer.pad_token_id
else:
pad_token_id = self.processor.pad_token_id

393
studio/backend/utils/datasets/dataset_utils.py
View file

@ -45,22 +45,21 @@ from .vlm_processing import generate_smart_vlm_instruction
from .data_collators import DeepSeekOCRDataCollator, VLMDataCollator
from .model_mappings import TEMPLATE_TO_MODEL_MAPPER
from loggers import get_logger
logger = get_logger(__name__)
def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"""
Lightweight format check without processing - for frontend validation.
Use this to quickly determine if user needs to manually map columns
before calling the full format_and_template_dataset().
Args:
dataset: HuggingFace dataset
is_vlm: Whether this is a Vision-Language Model dataset
Returns:
dict: {
"requires_manual_mapping": bool - True if user must map columns,
@ -71,8 +70,12 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"detected_text_column": str or None - For VLM only,
}
"""
columns = list(dataset.column_names) if hasattr(dataset, 'column_names') else list(next(iter(dataset)).keys())
columns = (
list(dataset.column_names)
if hasattr(dataset, "column_names")
else list(next(iter(dataset)).keys())
)
# Auto-detect multimodal data regardless of is_vlm flag
multimodal_info = detect_multimodal_dataset(dataset)
is_audio = multimodal_info.get("is_audio", False)
@ -185,11 +188,17 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
**audio_fields,
}
# Normalise any format-specific role to canonical chatml (user/assistant/system)
_TO_CHATML = {
"user": "user", "human": "user", "instruction": "user",
"assistant": "assistant", "gpt": "assistant", "output": "assistant",
"system": "system", "input": "system",
"user": "user",
"human": "user",
"instruction": "user",
"assistant": "assistant",
"gpt": "assistant",
"output": "assistant",
"system": "system",
"input": "system",
}
_CHATML_ROLE_ORDER = ("system", "user", "assistant")
_CHATML_TO_ALPACA = {"user": "instruction", "system": "input", "assistant": "output"}
@ -232,11 +241,21 @@ def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):
for col in role_groups[chatml_role]:
if col in examples:
content = examples[col][i]
convo.append({"role": chatml_role, "content": str(content) if content else ""})
convo.append(
{
"role": chatml_role,
"content": str(content) if content else "",
}
)
conversations.append(convo)
return {"conversations": conversations}
return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
return dataset.map(
_convert,
batched = True,
batch_size = batch_size,
remove_columns = dataset.column_names,
)
def _extract_column_value(val, col: str, label_mapping: dict) -> str:
@ -248,7 +267,7 @@ def _extract_column_value(val, col: str, label_mapping: dict) -> str:
inner = val["text"]
str_val = inner[0] if isinstance(inner, list) and inner else str(inner)
else:
str_val = json.dumps(val, ensure_ascii=False)
str_val = json.dumps(val, ensure_ascii = False)
elif isinstance(val, list):
str_val = val[0] if len(val) == 1 else ", ".join(str(v) for v in val)
else:
@ -286,6 +305,7 @@ def _apply_template_mapping(
role_groups[canonical].append(col)
import logging as _log
_log.getLogger(__name__).info(
f"Applying role mapping: sys={bool(system_prompt)}, "
f"user_cols={role_groups['user']}, asst_cols={role_groups['assistant']}, "
@ -326,8 +346,10 @@ def _apply_template_mapping(
return {"conversations": conversations}
return dataset.map(
_convert, batched=True, batch_size=batch_size,
remove_columns=dataset.column_names,
_convert,
batched = True,
batch_size = batch_size,
remove_columns = dataset.column_names,
)
@ -341,7 +363,11 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
Returns:
Dataset with instruction/input/output columns
"""
col_for: dict[str, str | None] = {"instruction": None, "input": None, "output": None}
col_for: dict[str, str | None] = {
"instruction": None,
"input": None,
"output": None,
}
for col_name, role in mapping.items():
canonical = _TO_CHATML.get(role)
alpaca_field = _CHATML_TO_ALPACA.get(canonical) if canonical else None
@ -352,25 +378,48 @@ def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
num = len(next(iter(examples.values())))
instructions, inputs, outputs = [], [], []
for i in range(num):
for field, dest in (("instruction", instructions), ("input", inputs), ("output", outputs)):
for field, dest in (
("instruction", instructions),
("input", inputs),
("output", outputs),
):
col = col_for[field]
val = str(examples[col][i]) if col and col in examples and examples[col][i] else ""
val = (
str(examples[col][i])
if col and col in examples and examples[col][i]
else ""
)
dest.append(val)
return {"instruction": instructions, "input": inputs, "output": outputs}
return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
return dataset.map(
_convert,
batched = True,
batch_size = batch_size,
remove_columns = dataset.column_names,
)
def format_dataset(
dataset,
format_type = "auto",
tokenizer = None,
aliases_for_system = ["system",],
aliases_for_user = ["user", "human", "input",],
aliases_for_assistant = ["gpt", "assistant", "output",],
batch_size = 1000,
num_proc = None,
auto_detect_custom = True,
format_type = "auto",
tokenizer = None,
aliases_for_system = [
"system",
],
aliases_for_user = [
"user",
"human",
"input",
],
aliases_for_assistant = [
"gpt",
"assistant",
"output",
],
batch_size = 1000,
num_proc = None,
auto_detect_custom = True,
custom_format_mapping = None,
):
"""
@ -395,13 +444,17 @@ def format_dataset(
if custom_format_mapping:
try:
if format_type == "alpaca":
mapped_dataset = _apply_user_mapping_alpaca(dataset, custom_format_mapping, batch_size)
mapped_dataset = _apply_user_mapping_alpaca(
dataset, custom_format_mapping, batch_size
)
final_format = "alpaca"
chat_column = None
else:
# auto / chatml / sharegpt / conversational — all produce chatml conversations
# (sharegpt is always standardized to role/content internally)
mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size)
mapped_dataset = _apply_user_mapping(
dataset, custom_format_mapping, batch_size
)
final_format = "chatml_conversations"
chat_column = "conversations"
@ -414,7 +467,9 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": [f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"]
"warnings": [
f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"
],
}
except Exception as e:
return {
@ -426,10 +481,9 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": [f"Failed to apply user mapping: {e}"]
"warnings": [f"Failed to apply user mapping: {e}"],
}
# Detect current format
detected = detect_dataset_format(dataset)
warnings = []
@ -442,7 +496,6 @@ def format_dataset(
# AUTO MODE: Keep format but standardize if needed
if format_type == "auto":
# Alpaca - keep as is
if detected["format"] == "alpaca":
return {
@ -454,16 +507,20 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
# ShareGPT - needs standardization
elif detected["format"] == "sharegpt":
try:
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
return {
"dataset": standardized,
@ -474,7 +531,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
except Exception as e:
warnings.append(f"Failed to standardize ShareGPT format: {e}")
@ -487,10 +544,14 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
elif detected["format"] == "chatml" and detected["chat_column"] in ["conversations", "messages", "texts"]:
elif detected["format"] == "chatml" and detected["chat_column"] in [
"conversations",
"messages",
"texts",
]:
return {
"dataset": dataset,
"detected_format": f"chatml_{detected['chat_column']}",
@ -500,13 +561,14 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
# Unknown - try standardization, if fails pass as is
else:
warnings.append(f"Unknown format detected. Keys found: {detected['sample_keys']}")
warnings.append(
f"Unknown format detected. Keys found: {detected['sample_keys']}"
)
# NEW: Try heuristic detection
if auto_detect_custom:
@ -514,7 +576,6 @@ def format_dataset(
if custom_mapping:
warnings.append(f"Auto-detected column mapping: {custom_mapping}")
def _apply_auto_mapping(examples):
conversations = []
num_examples = len(examples[list(examples.keys())[0]])
@ -523,25 +584,27 @@ def format_dataset(
all_columns = set(examples.keys())
mapped_columns = set(custom_mapping.keys())
preserved_columns = {
col: examples[col]
for col in all_columns - mapped_columns
col: examples[col] for col in all_columns - mapped_columns
}
for i in range(num_examples):
convo = []
for target_role in ['system', 'user', 'assistant']:
for target_role in ["system", "user", "assistant"]:
for col_name, role in custom_mapping.items():
if role == target_role and col_name in examples:
content = examples[col_name][i]
if content and str(content).strip():
convo.append({"role": role, "content": str(content)})
convo.append(
{"role": role, "content": str(content)}
)
conversations.append(convo)
return {"conversations": conversations, **preserved_columns}
try:
dataset = dataset.map(_apply_auto_mapping, batched=True, batch_size=batch_size)
dataset = dataset.map(
_apply_auto_mapping, batched = True, batch_size = batch_size
)
return {
"dataset": dataset,
"detected_format": "unknown",
@ -551,7 +614,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
except Exception as e:
warnings.append(f"Auto-detection failed: {e}")
@ -560,9 +623,13 @@ def format_dataset(
if detected["chat_column"]:
try:
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
warnings.append("Successfully standardized unknown format")
return {
@ -574,10 +641,12 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
except Exception as e:
warnings.append(f"Could not standardize: {e}. Passing dataset as-is.")
warnings.append(
f"Could not standardize: {e}. Passing dataset as-is."
)
# Return as-is with warnings
return {
@ -589,12 +658,11 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
# ALPACA MODE: Convert to Alpaca
elif format_type == "alpaca":
if detected["format"] == "alpaca":
return {
"dataset": dataset,
@ -605,16 +673,20 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
elif detected["format"] in ["sharegpt", "chatml"]:
# First standardize if ShareGPT
if detected["format"] == "sharegpt":
dataset = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
# Then convert to Alpaca
@ -628,7 +700,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
else:
@ -642,12 +714,11 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
# CHATML MODE: Convert to ChatML
elif format_type in ["chatml", "conversational", "sharegpt"]:
if detected["format"] == "alpaca":
converted = convert_alpaca_to_chatml(dataset, batch_size, num_proc)
return {
@ -659,14 +730,18 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
elif detected["format"] == "sharegpt":
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
return {
"dataset": standardized,
@ -677,7 +752,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
elif detected["format"] == "chatml":
@ -690,7 +765,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": []
"warnings": [],
}
else:
@ -698,9 +773,13 @@ def format_dataset(
if detected["chat_column"]:
try:
standardized = standardize_chat_format(
dataset, tokenizer, aliases_for_system,
aliases_for_user, aliases_for_assistant,
batch_size, num_proc
dataset,
tokenizer,
aliases_for_system,
aliases_for_user,
aliases_for_assistant,
batch_size,
num_proc,
)
return {
"dataset": standardized,
@ -711,7 +790,7 @@ def format_dataset(
"requires_manual_mapping": False,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
except Exception as e:
warnings.append(f"Standardization failed: {e}")
@ -725,7 +804,7 @@ def format_dataset(
"requires_manual_mapping": True,
"is_image": multimodal_info["is_image"],
"multimodal_info": multimodal_info,
"warnings": warnings
"warnings": warnings,
}
else:
@ -737,25 +816,34 @@ def format_and_template_dataset(
model_name,
tokenizer,
is_vlm = False,
format_type="auto",
format_type = "auto",
# VLM-specific parameters
vlm_instruction=None, # Now optional - will auto-generate
vlm_text_column=None,
vlm_image_column=None,
dataset_name=None,
custom_prompt_template=None,
add_eos_token=False,
remove_bos_prefix=False,
custom_format_mapping=None,
auto_detect_custom=True,
auto_detect_mapping=True,
aliases_for_system=["system",],
aliases_for_user=["user", "human", "input",],
aliases_for_assistant=["gpt", "assistant", "output",],
batch_size=1000,
num_proc=None,
progress_callback=None,
vlm_instruction = None, # Now optional - will auto-generate
vlm_text_column = None,
vlm_image_column = None,
dataset_name = None,
custom_prompt_template = None,
add_eos_token = False,
remove_bos_prefix = False,
custom_format_mapping = None,
auto_detect_custom = True,
auto_detect_mapping = True,
aliases_for_system = [
"system",
],
aliases_for_user = [
"user",
"human",
"input",
],
aliases_for_assistant = [
"gpt",
"assistant",
"output",
],
batch_size = 1000,
num_proc = None,
progress_callback = None,
):
"""
Convenience function that combines format_dataset and apply_chat_template_to_dataset.
@ -786,25 +874,27 @@ def format_and_template_dataset(
# Expect mapping like: {"image_col": "image", "caption_col": "text"}
user_vlm_image_column = None
user_vlm_text_column = None
for col, role in custom_format_mapping.items():
if role == "image":
user_vlm_image_column = col
elif role in ["text", "user", "caption", "assistant"]:
user_vlm_text_column = col
if user_vlm_image_column and user_vlm_text_column:
try:
dataset = convert_to_vlm_format(
dataset,
instruction=vlm_instruction,
text_column=user_vlm_text_column,
image_column=user_vlm_image_column,
dataset_name=dataset_name,
progress_callback=progress_callback,
instruction = vlm_instruction,
text_column = user_vlm_text_column,
image_column = user_vlm_image_column,
dataset_name = dataset_name,
progress_callback = progress_callback,
)
warnings.append(f"Applied user VLM mapping: image='{user_vlm_image_column}', text='{user_vlm_text_column}'")
warnings.append(
f"Applied user VLM mapping: image='{user_vlm_image_column}', text='{user_vlm_text_column}'"
)
return {
"dataset": dataset,
"detected_format": "user_mapped",
@ -826,7 +916,9 @@ def format_and_template_dataset(
f"text='{user_vlm_text_column}') failed: {e} — "
f"falling back to auto-detection"
)
logger.info(f"⚠️ User VLM mapping failed, falling back to auto-detection...")
logger.info(
f"⚠️ User VLM mapping failed, falling back to auto-detection..."
)
custom_format_mapping = None # clear so auto-detection runs below
else:
errors.append(
@ -850,10 +942,13 @@ def format_and_template_dataset(
if vlm_structure["format"] == "vlm_messages_llava":
try:
dataset = convert_llava_to_vlm_format(dataset)
warnings.append("Converted from Llava format (image indices) to standard VLM format")
warnings.append(
"Converted from Llava format (image indices) to standard VLM format"
)
except Exception as e:
errors.append(f"Failed to convert Llava format: {e}")
import traceback
traceback.print_exc()
return {
@ -872,15 +967,18 @@ def format_and_template_dataset(
try:
dataset = convert_sharegpt_with_images_to_vlm_format(
dataset,
image_column=vlm_structure["image_column"],
messages_column=vlm_structure["messages_column"],
dataset_name=dataset_name,
progress_callback=progress_callback,
image_column = vlm_structure["image_column"],
messages_column = vlm_structure["messages_column"],
dataset_name = dataset_name,
progress_callback = progress_callback,
)
warnings.append(
"Converted from ShareGPT+image format to standard VLM format"
)
warnings.append("Converted from ShareGPT+image format to standard VLM format")
except Exception as e:
errors.append(f"Failed to convert ShareGPT+image format: {e}")
import traceback
traceback.print_exc()
return {
@ -910,14 +1008,18 @@ def format_and_template_dataset(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=columns,
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = columns,
)
except Exception:
pass
errors.append(
friendly or f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
friendly
or f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
)
return {
"dataset": dataset,
@ -933,21 +1035,26 @@ def format_and_template_dataset(
try:
dataset = convert_to_vlm_format(
dataset,
instruction=vlm_instruction,
text_column=vlm_text_column,
image_column=vlm_image_column,
dataset_name=dataset_name,
progress_callback=progress_callback,
instruction = vlm_instruction,
text_column = vlm_text_column,
image_column = vlm_image_column,
dataset_name = dataset_name,
progress_callback = progress_callback,
)
if vlm_instruction:
warnings.append(f"Using user-provided instruction: '{vlm_instruction}'")
warnings.append(
f"Using user-provided instruction: '{vlm_instruction}'"
)
else:
warnings.append("Auto-generated instruction based on dataset analysis")
warnings.append(
"Auto-generated instruction based on dataset analysis"
)
except Exception as e:
errors.append(f"Failed to convert to VLM format: {e}")
import traceback
traceback.print_exc()
return {
@ -987,41 +1094,45 @@ def format_and_template_dataset(
# Step 1: Format the dataset
dataset_info = format_dataset(
dataset,
format_type=format_type,
tokenizer=tokenizer,
auto_detect_custom=auto_detect_custom,
custom_format_mapping=custom_format_mapping,
aliases_for_system=aliases_for_system,
aliases_for_user=aliases_for_user,
aliases_for_assistant=aliases_for_assistant,
batch_size=batch_size,
num_proc=num_proc,
format_type = format_type,
tokenizer = tokenizer,
auto_detect_custom = auto_detect_custom,
custom_format_mapping = custom_format_mapping,
aliases_for_system = aliases_for_system,
aliases_for_user = aliases_for_user,
aliases_for_assistant = aliases_for_assistant,
batch_size = batch_size,
num_proc = num_proc,
)
# Step 2: Apply chat template
# Gemma emits a leading <bos> that must be stripped for text-only chatml/sharegpt.
is_alpaca = format_type == "alpaca" or (format_type == "auto" and dataset_info["detected_format"] == "alpaca")
is_alpaca = format_type == "alpaca" or (
format_type == "auto" and dataset_info["detected_format"] == "alpaca"
)
is_gemma = "gemma" in model_name.lower()
if is_gemma and not dataset_info["is_image"] and not is_alpaca:
remove_bos_prefix = True
template_result = apply_chat_template_to_dataset(
dataset_info=dataset_info,
tokenizer=tokenizer,
model_name=model_name,
custom_prompt_template=custom_prompt_template,
add_eos_token=add_eos_token,
remove_bos_prefix=remove_bos_prefix,
custom_format_mapping=custom_format_mapping,
auto_detect_mapping=auto_detect_mapping,
batch_size=batch_size,
num_proc=num_proc,
dataset_info = dataset_info,
tokenizer = tokenizer,
model_name = model_name,
custom_prompt_template = custom_prompt_template,
add_eos_token = add_eos_token,
remove_bos_prefix = remove_bos_prefix,
custom_format_mapping = custom_format_mapping,
auto_detect_mapping = auto_detect_mapping,
batch_size = batch_size,
num_proc = num_proc,
)
# Step 3: Generate summary
summary = get_dataset_info_summary(dataset_info)
# Combine results
all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", [])
all_warnings = dataset_info.get("warnings", []) + template_result.get(
"warnings", []
)
all_errors = template_result.get("errors", [])
# If format_dataset returned "unknown" but apply_chat_template rescued

308
studio/backend/utils/datasets/format_conversion.py
View file

@ -12,19 +12,28 @@ import os
from datasets import IterableDataset
from loggers import get_logger
logger = get_logger(__name__)
def standardize_chat_format(
dataset,
tokenizer=None,
aliases_for_system=["system",],
aliases_for_user=["user", "human", "input",],
aliases_for_assistant=["gpt", "assistant", "output",],
batch_size=1000,
num_proc=None,
tokenizer = None,
aliases_for_system = [
"system",
],
aliases_for_user = [
"user",
"human",
"input",
],
aliases_for_assistant = [
"gpt",
"assistant",
"output",
],
batch_size = 1000,
num_proc = None,
):
"""
Our own standardization function that handles BOTH messages and conversations.
@ -67,22 +76,25 @@ def standardize_chat_format(
return dataset # Unexpected structure
keys = list(uniques.keys())
length_first = len(set(uniques[keys[0]]))
length_first = len(set(uniques[keys[0]]))
length_second = len(set(uniques[keys[1]]))
# Determine which is role and which is content
if length_first < length_second:
role_key = keys[0]
role_key = keys[0]
content_key = keys[1]
else:
role_key = keys[1]
role_key = keys[1]
content_key = keys[0]
# Mapping for aliases
aliases_mapping = {}
for x in aliases_for_system: aliases_mapping[x] = "system"
for x in aliases_for_user: aliases_mapping[x] = "user"
for x in aliases_for_assistant: aliases_mapping[x] = "assistant"
for x in aliases_for_system:
aliases_mapping[x] = "system"
for x in aliases_for_user:
aliases_mapping[x] = "user"
for x in aliases_for_assistant:
aliases_mapping[x] = "assistant"
def _standardize_dataset(examples):
convos = examples[chat_column]
@ -109,10 +121,9 @@ def standardize_chat_format(
return {chat_column: all_convos}
dataset_map_kwargs = {
'batched': True,
'batch_size': batch_size,
"batched": True,
"batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@ -123,13 +134,13 @@ def standardize_chat_format(
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Standardizing chat format"
dataset_map_kwargs["num_proc"] = num_proc
dataset_map_kwargs["desc"] = "Standardizing chat format"
return dataset.map(_standardize_dataset, **dataset_map_kwargs)
def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
def convert_chatml_to_alpaca(dataset, batch_size = 1000, num_proc = None):
"""
Converts ChatML format (messages OR conversations) to Alpaca format.
Handles both standardized and ShareGPT formats.
@ -142,10 +153,16 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
def _convert(examples):
# Auto-detect which column name is used
chatml_data = examples.get("messages") or examples.get("conversations") or examples.get("texts")
chatml_data = (
examples.get("messages")
or examples.get("conversations")
or examples.get("texts")
)
if chatml_data is None:
raise ValueError("No 'messages' or 'conversations' or 'texts' column found.")
raise ValueError(
"No 'messages' or 'conversations' or 'texts' column found."
)
instructions = []
outputs = []
@ -172,15 +189,11 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
inputs.append("") # Alpaca typically has empty input
outputs.append(output)
return {
"instruction": instructions,
"input": inputs,
"output": outputs
}
return {"instruction": instructions, "input": inputs, "output": outputs}
dataset_map_kwargs = {
'batched': True,
'batch_size': batch_size,
"batched": True,
"batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@ -191,13 +204,13 @@ def convert_chatml_to_alpaca(dataset, batch_size=1000, num_proc=None):
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting ChatML to Alpaca format"
dataset_map_kwargs["num_proc"] = num_proc
dataset_map_kwargs["desc"] = "Converting ChatML to Alpaca format"
return dataset.map(_convert, **dataset_map_kwargs)
def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
def convert_alpaca_to_chatml(dataset, batch_size = 1000, num_proc = None):
"""
Converts Alpaca format to ChatML format.
@ -222,15 +235,15 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
# Build conversation in standard ChatML format
convo = [
{"role": "user", "content": user_content},
{"role": "assistant", "content": output}
{"role": "assistant", "content": output},
]
conversations.append(convo)
return {"conversations": conversations}
dataset_map_kwargs = {
'batched': True,
'batch_size': batch_size,
"batched": True,
"batch_size": batch_size,
}
if not isinstance(dataset, IterableDataset):
@ -241,8 +254,8 @@ def convert_alpaca_to_chatml(dataset, batch_size=1000, num_proc=None):
else:
num_proc = safe_num_proc(num_proc)
dataset_map_kwargs['num_proc'] = num_proc
dataset_map_kwargs['desc'] = "Converting Alpaca to ChatML format"
dataset_map_kwargs["num_proc"] = num_proc
dataset_map_kwargs["desc"] = "Converting Alpaca to ChatML format"
return dataset.map(_convert, **dataset_map_kwargs)
@ -262,11 +275,11 @@ def _format_eta(seconds):
def convert_to_vlm_format(
dataset,
instruction=None,
text_column="text",
image_column="image",
dataset_name=None,
progress_callback=None,
instruction = None,
text_column = "text",
image_column = "image",
dataset_name = None,
progress_callback = None,
):
"""
Converts simple {image, text} format to VLM messages format.
@ -290,27 +303,31 @@ def convert_to_vlm_format(
def _notify(msg):
"""Send status update to the training overlay if callback is available."""
if progress_callback:
progress_callback(status_message=msg)
progress_callback(status_message = msg)
# Generate smart instruction if not provided
if instruction is None:
instruction_info = generate_smart_vlm_instruction(
dataset,
text_column=text_column,
image_column=image_column,
dataset_name=dataset_name,
text_column = text_column,
image_column = image_column,
dataset_name = dataset_name,
)
instruction = instruction_info["instruction"]
instruction_column = instruction_info.get("instruction_column")
uses_dynamic = instruction_info["uses_dynamic_instruction"]
logger.info(f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}")
logger.info(
f"📝 Auto-detected instruction type: {instruction_info['instruction_type']}"
)
logger.info(f"📝 Confidence: {instruction_info['confidence']:.2f}")
if not uses_dynamic:
logger.info(f"📝 Using instruction: '{instruction}'")
else:
logger.info(f"📝 Using dynamic instructions from column: '{instruction_column}'")
logger.info(
f"📝 Using dynamic instructions from column: '{instruction_column}'"
)
else:
instruction_column = None
uses_dynamic = False
@ -324,13 +341,17 @@ def convert_to_vlm_format(
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
with fsspec.open(image_data, "rb", expand=True) as f:
with fsspec.open(image_data, "rb", expand = True) as f:
image_data = Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
# Bare filename → resolve via HF repo lookup
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
dataset_name, _image_lookup[image_data], repo_type="dataset",
dataset_name,
_image_lookup[image_data],
repo_type = "dataset",
)
image_data = Image.open(local_path).convert("RGB")
else:
@ -340,6 +361,7 @@ def convert_to_vlm_format(
text_data = sample[text_column]
if isinstance(text_data, list) and len(text_data) > 0:
import random
text_data = random.choice(text_data)
# Get instruction (static or dynamic)
@ -354,15 +376,10 @@ def convert_to_vlm_format(
"role": "user",
"content": [
{"type": "text", "text": current_instruction},
{"type": "image", "image": image_data} # PIL object
]
{"type": "image", "image": image_data}, # PIL object
],
},
{
"role": "assistant",
"content": [
{"type": "text", "text": text_data}
]
}
{"role": "assistant", "content": [{"type": "text", "text": text_data}]},
]
# Return dict with messages
@ -370,13 +387,15 @@ def convert_to_vlm_format(
total = len(dataset)
first_image = next(iter(dataset))[image_column]
has_urls = isinstance(first_image, str) and first_image.startswith(("http://", "https://"))
has_urls = isinstance(first_image, str) and first_image.startswith(
("http://", "https://")
)
# ── Bare-filename detection: images stored as filenames (e.g. "img_001.png")
# that don't exist locally. Build a basename→repo_path lookup so we can
# resolve them via hf_hub_download during conversion.
_image_lookup = None
_IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff')
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff")
if (
not has_urls
and isinstance(first_image, str)
@ -385,18 +404,25 @@ def convert_to_vlm_format(
):
try:
from huggingface_hub import HfApi
_notify("Resolving image filenames from HF repo...")
logger.info(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
logger.info(
f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup..."
)
repo_files = HfApi().list_repo_files(dataset_name, repo_type = "dataset")
_image_lookup = {
os.path.basename(f): f
for f in repo_files
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS)
}
if first_image in _image_lookup:
logger.info(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
logger.info(
f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')"
)
else:
logger.info(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
logger.info(
f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open"
)
_image_lookup = None
except Exception as e:
logger.info(f"⚠️ Failed to build HF repo image lookup: {e}")
@ -413,15 +439,19 @@ def convert_to_vlm_format(
num_workers = safe_num_proc()
_notify(f"Probing {PROBE_SIZE} image URLs with {num_workers} workers...")
logger.info(f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers...")
logger.info(
f"🔍 Probing {PROBE_SIZE}/{total} image URLs with {num_workers} workers..."
)
probe_samples = [dataset[i] for i in range(PROBE_SIZE)]
probe_ok = 0
probe_fail = 0
probe_start = time.time()
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = {executor.submit(_convert_single_sample, s): s for s in probe_samples}
with ThreadPoolExecutor(max_workers = num_workers) as executor:
futures = {
executor.submit(_convert_single_sample, s): s for s in probe_samples
}
for future in as_completed(futures):
try:
future.result()
@ -443,9 +473,12 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=[image_column, text_column],
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = [image_column, text_column],
)
except Exception:
pass
@ -471,7 +504,9 @@ def convert_to_vlm_format(
if probe_fail > 0:
info_msg += f" | {fail_rate:.0%} broken URLs will be skipped"
logger.info(f"✅ Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s")
logger.info(
f"✅ Probe passed: {probe_ok}/{probe_total} ok, {probe_fail} failed ({fail_rate:.0%}), {throughput:.1f} img/s"
)
logger.info(f"⏱️ Estimated time for {total:,} samples: ~{eta_str}")
_notify(info_msg)
@ -496,8 +531,11 @@ def convert_to_vlm_format(
batch_end = min(batch_start + batch_size, total)
batch_samples = [dataset[i] for i in range(batch_start, batch_end)]
with ThreadPoolExecutor(max_workers=num_workers) as executor:
futures = {executor.submit(_convert_single_sample, s): i for i, s in enumerate(batch_samples)}
with ThreadPoolExecutor(max_workers = num_workers) as executor:
futures = {
executor.submit(_convert_single_sample, s): i
for i, s in enumerate(batch_samples)
}
batch_results = [None] * len(batch_samples)
for future in as_completed(futures):
idx = futures[future]
@ -506,9 +544,13 @@ def convert_to_vlm_format(
except Exception as e:
failed_count += 1
if failed_count == 1:
print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
print(
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
)
if failed_count == 1:
logger.info(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
logger.info(
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
)
converted_list.extend(r for r in batch_results if r is not None)
@ -519,11 +561,13 @@ def convert_to_vlm_format(
remaining_time = (total - done) / rate if rate > 0 else 0
eta_str = _format_eta(remaining_time)
progress_msg = f"Downloading images: {done:,}/{total:,} ({done*100//total}%) | ~{eta_str} remaining | {failed_count} skipped"
logger.info(f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}")
logger.info(
f" [{done}/{total}] {rate:.1f} img/s, {failed_count} failed, ETA {eta_str}"
)
_notify(progress_msg)
else:
# Sequential conversion for local/embedded images (fast, no I/O bottleneck)
pbar = tqdm(dataset, total=total, desc="Converting VLM samples", unit="sample")
pbar = tqdm(dataset, total = total, desc = "Converting VLM samples", unit = "sample")
for sample in pbar:
try:
converted_list.append(_convert_single_sample(sample))
@ -534,13 +578,17 @@ def convert_to_vlm_format(
print(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
if failed_count == 1:
# Log the first failure to aid debugging
logger.info(f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}")
pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
logger.info(
f"⚠️ First VLM conversion failure: {type(e).__name__}: {e}"
)
pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
pbar.close()
if failed_count > 0:
fail_rate = failed_count / total
logger.info(f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images")
logger.info(
f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images"
)
# For datasets that skipped the probe (small URL datasets), check fail rate now
if has_urls and fail_rate >= MAX_FAIL_RATE:
issues = [
@ -550,9 +598,12 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=[image_column, text_column],
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = [image_column, text_column],
)
except Exception:
pass
@ -573,14 +624,18 @@ def convert_to_vlm_format(
friendly = None
try:
from .llm_assist import llm_generate_dataset_warning
friendly = llm_generate_dataset_warning(
issues, dataset_name=dataset_name, modality="vision",
column_names=[image_column, text_column],
issues,
dataset_name = dataset_name,
modality = "vision",
column_names = [image_column, text_column],
)
except Exception:
pass
raise ValueError(
friendly or (
friendly
or (
f"All {total} samples failed during VLM conversion — no usable images found. "
"This dataset may contain only image URLs that are no longer accessible."
)
@ -595,10 +650,10 @@ def convert_to_vlm_format(
def convert_sharegpt_with_images_to_vlm_format(
dataset,
image_column="image",
messages_column="conversations",
dataset_name=None,
progress_callback=None,
image_column = "image",
messages_column = "conversations",
dataset_name = None,
progress_callback = None,
):
"""
Converts ShareGPT/ChatML datasets that have a separate image column and
@ -619,16 +674,18 @@ def convert_sharegpt_with_images_to_vlm_format(
from PIL import Image
from tqdm import tqdm
_IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff')
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff")
_ROLE_MAP = {
"human": "user", "user": "user",
"gpt": "assistant", "assistant": "assistant",
"human": "user",
"user": "user",
"gpt": "assistant",
"assistant": "assistant",
"system": "system",
}
def _notify(msg):
if progress_callback:
progress_callback(status_message=msg)
progress_callback(status_message = msg)
# ── Resolve image loading strategy (same 3-tier as convert_to_vlm_format) ──
total = len(dataset)
@ -643,9 +700,12 @@ def convert_sharegpt_with_images_to_vlm_format(
):
try:
from huggingface_hub import HfApi
_notify("Resolving image filenames from HF repo...")
logger.info(f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup...")
repo_files = HfApi().list_repo_files(dataset_name, repo_type="dataset")
logger.info(
f"🔍 Image column contains bare filenames (e.g. '{first_image}') — building repo lookup..."
)
repo_files = HfApi().list_repo_files(dataset_name, repo_type = "dataset")
_image_lookup = {
os.path.basename(f): f
for f in repo_files
@ -656,9 +716,13 @@ def convert_sharegpt_with_images_to_vlm_format(
if any(f.lower().endswith(ext) for ext in _IMAGE_EXTS):
_image_lookup[f] = f
if first_image in _image_lookup:
logger.info(f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')")
logger.info(
f"✅ Matched {len(_image_lookup)} image files in repo (e.g. '{first_image}' → '{_image_lookup[first_image]}')"
)
else:
logger.info(f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open")
logger.info(
f"⚠️ Built lookup with {len(_image_lookup)} images but '{first_image}' not found — falling back to local open"
)
_image_lookup = None
except Exception as e:
logger.info(f"⚠️ Failed to build HF repo image lookup: {e}")
@ -666,25 +730,32 @@ def convert_sharegpt_with_images_to_vlm_format(
def _resolve_image(image_data):
"""Resolve image data to a PIL Image object."""
if hasattr(image_data, 'size') and hasattr(image_data, 'mode'):
if hasattr(image_data, "size") and hasattr(image_data, "mode"):
return image_data # Already PIL
if isinstance(image_data, str):
if image_data.startswith(("http://", "https://")):
import fsspec
from io import BytesIO
with fsspec.open(image_data, "rb", expand=True) as f:
with fsspec.open(image_data, "rb", expand = True) as f:
return Image.open(BytesIO(f.read())).convert("RGB")
elif _image_lookup is not None and image_data in _image_lookup:
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
dataset_name, _image_lookup[image_data], repo_type="dataset",
dataset_name,
_image_lookup[image_data],
repo_type = "dataset",
)
return Image.open(local_path).convert("RGB")
else:
return Image.open(image_data).convert("RGB")
if isinstance(image_data, dict) and ("bytes" in image_data or "path" in image_data):
if isinstance(image_data, dict) and (
"bytes" in image_data or "path" in image_data
):
if image_data.get("bytes"):
from io import BytesIO
return Image.open(BytesIO(image_data["bytes"])).convert("RGB")
if image_data.get("path"):
return Image.open(image_data["path"]).convert("RGB")
@ -726,7 +797,7 @@ def convert_sharegpt_with_images_to_vlm_format(
converted_list = []
failed_count = 0
pbar = tqdm(dataset, total=total, desc="Converting ShareGPT+image", unit="sample")
pbar = tqdm(dataset, total = total, desc = "Converting ShareGPT+image", unit = "sample")
for sample in pbar:
try:
converted_list.append(_convert_single_sample(sample))
@ -734,11 +805,13 @@ def convert_sharegpt_with_images_to_vlm_format(
failed_count += 1
if failed_count == 1:
logger.info(f"⚠️ First conversion failure: {type(e).__name__}: {e}")
pbar.set_postfix(ok=len(converted_list), failed=failed_count, refresh=False)
pbar.set_postfix(ok = len(converted_list), failed = failed_count, refresh = False)
pbar.close()
if failed_count > 0:
logger.info(f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples")
logger.info(
f"⚠️ Skipped {failed_count}/{total} ({failed_count*100//total}%) samples"
)
if len(converted_list) == 0:
raise ValueError(
@ -764,7 +837,9 @@ def convert_llava_to_vlm_format(dataset):
"""
from PIL import Image
logger.info(f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format...")
logger.info(
f"🔄 Converting {len(dataset)} samples from Llava format to standard VLM format..."
)
def _convert_single_sample(sample):
"""Convert a single llava sample to standard VLM format."""
@ -787,10 +862,12 @@ def convert_llava_to_vlm_format(dataset):
if isinstance(pil_image, str):
pil_image = Image.open(pil_image).convert("RGB")
new_content.append({
"type": "image",
"image": pil_image # Actual PIL object
})
new_content.append(
{
"type": "image",
"image": pil_image, # Actual PIL object
}
)
else:
# No index, try to use first image
if len(images) > 0:
@ -798,22 +875,13 @@ def convert_llava_to_vlm_format(dataset):
if isinstance(pil_image, str):
pil_image = Image.open(pil_image).convert("RGB")
new_content.append({
"type": "image",
"image": pil_image
})
new_content.append({"type": "image", "image": pil_image})
elif item["type"] == "text":
# Keep text as-is (only type + text)
new_content.append({
"type": "text",
"text": item.get("text", "")
})
new_content.append({"type": "text", "text": item.get("text", "")})
new_messages.append({
"role": msg["role"],
"content": new_content
})
new_messages.append({"role": msg["role"], "content": new_content})
return {"messages": new_messages}

300
studio/backend/utils/datasets/format_detection.py
View file

@ -13,7 +13,10 @@ import re
def _keyword_in_column(keyword: str, col_name: str) -> bool:
"""Word-boundary keyword match to avoid false positives like 'pic' in 'topic'."""
return re.search(r'\b' + re.escape(keyword) + r'\b', col_name, re.IGNORECASE) is not None
return (
re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE)
is not None
)
def detect_dataset_format(dataset):
@ -37,7 +40,7 @@ def detect_dataset_format(dataset):
"format": "alpaca",
"chat_column": None,
"needs_standardization": False,
"sample_keys": []
"sample_keys": [],
}
# Check for chat-based formats (messages or conversations)
@ -65,7 +68,7 @@ def detect_dataset_format(dataset):
"format": "sharegpt",
"chat_column": chat_column,
"needs_standardization": True,
"sample_keys": list(msg_keys)
"sample_keys": list(msg_keys),
}
# ChatML uses "role" and "content"
@ -74,7 +77,7 @@ def detect_dataset_format(dataset):
"format": "chatml",
"chat_column": chat_column,
"needs_standardization": False,
"sample_keys": list(msg_keys)
"sample_keys": list(msg_keys),
}
# Unknown structure but has chat column
@ -83,7 +86,7 @@ def detect_dataset_format(dataset):
"format": "unknown",
"chat_column": chat_column,
"needs_standardization": None,
"sample_keys": list(msg_keys)
"sample_keys": list(msg_keys),
}
except Exception as e:
return {
@ -91,7 +94,7 @@ def detect_dataset_format(dataset):
"chat_column": chat_column,
"needs_standardization": None,
"sample_keys": [],
"error": str(e)
"error": str(e),
}
# No recognized format
@ -99,7 +102,7 @@ def detect_dataset_format(dataset):
"format": "unknown",
"chat_column": None,
"needs_standardization": None,
"sample_keys": []
"sample_keys": [],
}
@ -120,49 +123,86 @@ def detect_custom_format_heuristic(dataset):
# Keywords
assistant_words = [
'output', 'answer', 'response', 'assistant', 'completion',
'expected', 'recommendation', 'reply', 'result', 'target',
'solution', 'explanation', 'solve'
"output",
"answer",
"response",
"assistant",
"completion",
"expected",
"recommendation",
"reply",
"result",
"target",
"solution",
"explanation",
"solve",
]
# Split into high/low priority
user_words_high_priority = [
'input', 'question', 'query', 'prompt', 'instruction',
'request', 'snippet', 'user', 'text',
'problem', 'exercise'
"input",
"question",
"query",
"prompt",
"instruction",
"request",
"snippet",
"user",
"text",
"problem",
"exercise",
]
user_words_low_priority = ['task'] # Ambiguous - can be user OR system
user_words_low_priority = ["task"] # Ambiguous - can be user OR system
user_words = user_words_high_priority + user_words_low_priority
system_words = [
'system', 'context', 'description', 'persona', 'role',
'template', 'task' # Also in system
"system",
"context",
"description",
"persona",
"role",
"template",
"task", # Also in system
]
# Metadata columns to ignore
metadata_exact_match = {
'id', 'idx', 'index', 'key', 'timestamp', 'date',
'metadata', 'source', 'kind', 'type', 'category',
'score', 'label', 'tag', 'inference_mode'
"id",
"idx",
"index",
"key",
"timestamp",
"date",
"metadata",
"source",
"kind",
"type",
"category",
"score",
"label",
"tag",
"inference_mode",
}
metadata_prefix_patterns = [
'problem_type', 'problem_source',
'generation_model', 'pass_rate',
"problem_type",
"problem_source",
"generation_model",
"pass_rate",
]
priority_patterns = {
'generated': 100,
'gen_': 90,
'model_': 80,
'predicted': 70,
'completion': 60,
"generated": 100,
"gen_": 90,
"model_": 80,
"predicted": 70,
"completion": 60,
}
def has_keyword(col_name, keywords):
"""Check if any keyword appears in column name."""
col_lower = col_name.lower()
col_normalized = col_lower.replace('_', '').replace('-', '').replace(' ', '')
col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "")
for keyword in keywords:
if keyword in col_lower or keyword in col_normalized:
@ -180,13 +220,16 @@ def detect_custom_format_heuristic(dataset):
return True
for pattern in metadata_prefix_patterns:
if col_lower.startswith(pattern.split('_')[0] + '_') and col_lower != pattern:
if '_' in col_lower:
prefix = col_lower.split('_')[0]
if prefix in ['generation', 'pass', 'inference']:
if (
col_lower.startswith(pattern.split("_")[0] + "_")
and col_lower != pattern
):
if "_" in col_lower:
prefix = col_lower.split("_")[0]
if prefix in ["generation", "pass", "inference"]:
return True
if len(col_lower) <= 2 and not col_lower in ['qa', 'q', 'a']:
if len(col_lower) <= 2 and not col_lower in ["qa", "q", "a"]:
return True
return False
@ -221,16 +264,18 @@ def detect_custom_format_heuristic(dataset):
score += 10
# Penalize ambiguous keywords when scoring for user
if role_type == 'user':
if role_type == "user":
col_lower = col_name.lower()
# If column is ONLY "task" (or task_xxx), give it lower priority for user role
if 'task' in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
if "task" in col_lower and not any(
kw in col_lower for kw in user_words_high_priority
):
score -= 15 # Significant penalty so other user columns win
priority_bonus = get_priority_score(col_name)
score += priority_bonus
if role_type in ['assistant', 'user']:
if role_type in ["assistant", "user"]:
avg_length = get_content_length(col_name)
if num_candidates > 1:
@ -256,20 +301,24 @@ def detect_custom_format_heuristic(dataset):
content_columns = [col for col in all_columns if not is_metadata(col)]
# Count candidates first
assistant_potential = [col for col in content_columns if has_keyword(col, assistant_words)]
assistant_potential = [
col for col in content_columns if has_keyword(col, assistant_words)
]
user_potential = [col for col in content_columns if has_keyword(col, user_words)]
# STEP 1: Find best ASSISTANT column
assistant_candidates = []
for col in assistant_potential:
score = score_column(col, assistant_words, 'assistant', len(assistant_potential))
score = score_column(
col, assistant_words, "assistant", len(assistant_potential)
)
if score > 0:
assistant_candidates.append((col, score))
if assistant_candidates:
assistant_candidates.sort(key=lambda x: x[1], reverse=True)
assistant_candidates.sort(key = lambda x: x[1], reverse = True)
assistant_col = assistant_candidates[0][0]
mapping[assistant_col] = 'assistant'
mapping[assistant_col] = "assistant"
else:
assistant_col = None
@ -278,14 +327,14 @@ def detect_custom_format_heuristic(dataset):
for col in user_potential:
if col == assistant_col:
continue
score = score_column(col, user_words, 'user', len(user_potential))
score = score_column(col, user_words, "user", len(user_potential))
if score > 0:
user_candidates.append((col, score))
if user_candidates:
user_candidates.sort(key=lambda x: x[1], reverse=True)
user_candidates.sort(key = lambda x: x[1], reverse = True)
user_col = user_candidates[0][0]
mapping[user_col] = 'user'
mapping[user_col] = "user"
else:
user_col = None
@ -296,7 +345,7 @@ def detect_custom_format_heuristic(dataset):
for col in remaining_columns:
if has_keyword(col, system_words):
# Found a system match in remaining columns
mapping[col] = 'system'
mapping[col] = "system"
system_col = col
break
@ -309,22 +358,22 @@ def detect_custom_format_heuristic(dataset):
# If no strong keyword match, decide based on what's missing
if not has_keyword(remaining_col, user_words + assistant_words):
mapping[remaining_col] = 'system'
mapping[remaining_col] = "system"
elif user_col is None:
# No user column yet, assign this as user
mapping[remaining_col] = 'user'
mapping[remaining_col] = "user"
else:
# Already have user + assistant, treat as system context
mapping[remaining_col] = 'system'
mapping[remaining_col] = "system"
# VALIDATION: Ensure we have at least user + assistant
has_user = any(role == 'user' for role in mapping.values())
has_assistant = any(role == 'assistant' for role in mapping.values())
has_user = any(role == "user" for role in mapping.values())
has_assistant = any(role == "assistant" for role in mapping.values())
if not has_user and len(remaining_columns) > 0:
for col in remaining_columns:
if col not in mapping:
mapping[col] = 'user'
mapping[col] = "user"
has_user = True
break
@ -358,14 +407,27 @@ def detect_multimodal_dataset(dataset):
# Keywords that indicate image data
image_keywords = [
'image', 'img', 'pixel',
'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg',
'photo', 'pic', 'picture', 'visual',
'file_name', 'filename',
"image",
"img",
"pixel",
"jpg",
"jpeg",
"png",
"webp",
"bmp",
"gif",
"tiff",
"svg",
"photo",
"pic",
"picture",
"visual",
"file_name",
"filename",
]
# Keywords that indicate audio data
audio_keywords = ['audio', 'speech', 'wav', 'waveform', 'sound']
audio_keywords = ["audio", "speech", "wav", "waveform", "sound"]
multimodal_columns = []
audio_columns = []
@ -419,7 +481,7 @@ def detect_multimodal_dataset(dataset):
# Detect text column for audio datasets
detected_text_col = None
if audio_columns:
text_keywords = ['text', 'sentence', 'transcript', 'transcription', 'label']
text_keywords = ["text", "sentence", "transcript", "transcription", "label"]
for col_name in column_names:
if col_name.lower() in text_keywords:
detected_text_col = col_name
@ -430,7 +492,7 @@ def detect_multimodal_dataset(dataset):
# Detect speaker_id column for TTS datasets (CSM, Orpheus, Spark)
detected_speaker_col = None
if audio_columns:
speaker_keywords = ['source', 'speaker', 'speaker_id']
speaker_keywords = ["source", "speaker", "speaker_id"]
for col_name in column_names:
if col_name.lower() in speaker_keywords:
detected_speaker_col = col_name
@ -456,6 +518,7 @@ def _is_image_value(value) -> bool:
# PIL Image instance
try:
from PIL.Image import Image as PILImage
if isinstance(value, PILImage):
return True
except ImportError:
@ -470,7 +533,9 @@ def _is_image_value(value) -> bool:
if "bytes" in value and "path" in value:
# Check path extension to exclude audio files
path = value.get("path") or ""
if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
if isinstance(path, str) and any(
path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
):
return False
return True
@ -479,11 +544,13 @@ def _is_image_value(value) -> bool:
return _has_image_header(value)
# String that looks like an image file path or URL
_IMAGE_EXTS = ('.png', '.jpg', '.jpeg', '.webp', '.gif', '.bmp', '.tiff', '.svg')
_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg")
if isinstance(value, str) and len(value) < 1000:
lower = value.strip().lower()
# Image URL (http://... ending in image extension)
if lower.startswith(("http://", "https://")) and any(lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS):
if lower.startswith(("http://", "https://")) and any(
lower.split("?")[0].endswith(ext) for ext in _IMAGE_EXTS
):
return True
# Image file path (relative or absolute path ending in image extension)
if any(lower.endswith(ext) for ext in _IMAGE_EXTS):
@ -493,7 +560,15 @@ def _is_image_value(value) -> bool:
_AUDIO_EXTENSIONS = (
".wav", ".mp3", ".flac", ".ogg", ".opus", ".m4a", ".aac", ".wma", ".webm",
".wav",
".mp3",
".flac",
".ogg",
".opus",
".m4a",
".aac",
".wma",
".webm",
)
@ -509,7 +584,9 @@ def _is_audio_value(value) -> bool:
# Undecoded/streaming → {"bytes": b"...", "path": "some.wav"}
if "bytes" in value or "path" in value:
path = value.get("path") or ""
if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
if isinstance(path, str) and any(
path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
):
return True
return False
@ -520,19 +597,19 @@ def _has_image_header(data: bytes) -> bool:
if len(data) < 4:
return False
# JPEG
if data[:2] == b'\xff\xd8':
if data[:2] == b"\xff\xd8":
return True
# PNG
if data[:4] == b'\x89PNG':
if data[:4] == b"\x89PNG":
return True
# GIF
if data[:3] == b'GIF':
if data[:3] == b"GIF":
return True
# WebP
if data[:4] == b'RIFF' and len(data) >= 12 and data[8:12] == b'WEBP':
if data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP":
return True
# BMP
if data[:2] == b'BM':
if data[:2] == b"BM":
return True
return False
@ -568,10 +645,13 @@ def detect_vlm_dataset_structure(dataset):
if isinstance(content, list) and len(content) > 0:
if isinstance(content[0], dict) and "type" in content[0]:
# Check for llava format
has_index = any('index' in item for item in content if isinstance(item, dict))
has_images_column = 'images' in column_names
has_index = any(
"index" in item
for item in content
if isinstance(item, dict)
)
has_images_column = "images" in column_names
if has_index and has_images_column:
return {
@ -583,7 +663,11 @@ def detect_vlm_dataset_structure(dataset):
}
# Standard VLM format
has_image = any('image' in item for item in content if isinstance(item, dict))
has_image = any(
"image" in item
for item in content
if isinstance(item, dict)
)
if has_image:
return {
"format": "vlm_messages",
@ -637,26 +721,65 @@ def detect_vlm_dataset_structure(dataset):
# Define metadata patterns to EXCLUDE
metadata_patterns = {
'suffixes': ['_id', '_url', '_name', '_filename', '_uri', '_link', '_key', '_index'],
'prefixes': ['id_', 'url_', 'name_', 'filename_', 'uri_', 'link_', 'key_', 'index_'],
"suffixes": [
"_id",
"_url",
"_name",
"_filename",
"_uri",
"_link",
"_key",
"_index",
],
"prefixes": [
"id_",
"url_",
"name_",
"filename_",
"uri_",
"link_",
"key_",
"index_",
],
}
# Image-related keywords
image_keywords = ['image', 'img', 'photo', 'picture', 'pic', 'visual', 'scan', 'file_name', 'filename']
image_keywords = [
"image",
"img",
"photo",
"picture",
"pic",
"visual",
"scan",
"file_name",
"filename",
]
# Text-related keywords
text_keywords = ['text', 'caption', 'captions', 'description', 'answer', 'output', 'response', 'label']
text_keywords = [
"text",
"caption",
"captions",
"description",
"answer",
"output",
"response",
"label",
]
def is_metadata_column(col_name):
"""Check if column name looks like metadata."""
col_lower = col_name.lower()
# Check suffixes
if any(col_lower.endswith(suffix) for suffix in metadata_patterns['suffixes']):
if any(col_lower.endswith(suffix) for suffix in metadata_patterns["suffixes"]):
return True
# Check prefixes
if any(col_lower.startswith(prefix) for prefix in metadata_patterns['prefixes']):
if any(
col_lower.startswith(prefix) for prefix in metadata_patterns["prefixes"]
):
return True
return False
@ -664,11 +787,13 @@ def detect_vlm_dataset_structure(dataset):
def _score_image_candidate(col, sample_value):
"""Score a candidate image column by how resolvable its value is."""
# PIL Image object (highest priority - already loaded)
if hasattr(sample_value, 'size') and hasattr(sample_value, 'mode'):
if hasattr(sample_value, "size") and hasattr(sample_value, "mode"):
return 100
# Dict with image data (bytes/path from HF Image feature)
if isinstance(sample_value, dict) and ('bytes' in sample_value or 'path' in sample_value):
if isinstance(sample_value, dict) and (
"bytes" in sample_value or "path" in sample_value
):
return 75
if isinstance(sample_value, str):
@ -693,13 +818,16 @@ def detect_vlm_dataset_structure(dataset):
# Local file — check it exists
if not sample_value.startswith(("http://", "https://")):
return os.path.exists(sample_value) # bare filenames return False here, that's OK
return os.path.exists(
sample_value
) # bare filenames return False here, that's OK
# URL — quick HEAD request with short timeout
try:
import urllib.request
req = urllib.request.Request(sample_value, method="HEAD")
resp = urllib.request.urlopen(req, timeout=3)
req = urllib.request.Request(sample_value, method = "HEAD")
resp = urllib.request.urlopen(req, timeout = 3)
return resp.status < 400
except Exception:
return False
@ -732,7 +860,7 @@ def detect_vlm_dataset_structure(dataset):
if not candidates:
return None
candidates.sort(key=lambda x: x[1], reverse=True)
candidates.sort(key = lambda x: x[1], reverse = True)
# Single candidate or top candidate is PIL/dict — no probing needed
if len(candidates) == 1 or candidates[0][1] >= 75:
@ -766,14 +894,18 @@ def detect_vlm_dataset_structure(dataset):
# Longer text = higher priority (likely content, not just a label)
priority = min(len(sample_value), 1000) # Cap at 1000
candidates.append((col, priority))
elif isinstance(sample_value, list) and len(sample_value) > 0 and isinstance(sample_value[0], str):
elif (
isinstance(sample_value, list)
and len(sample_value) > 0
and isinstance(sample_value[0], str)
):
# List of strings (e.g. captions list) — lower priority than plain strings
priority = min(len(sample_value[0]), 1000) // 2
candidates.append((col, priority))
# Return highest priority candidate
if candidates:
candidates.sort(key=lambda x: x[1], reverse=True)
candidates.sort(key = lambda x: x[1], reverse = True)
return candidates[0][0]
return None

207
studio/backend/utils/datasets/llm_assist.py
View file

@ -42,33 +42,45 @@ def precache_helper_gguf():
return
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
try:
from huggingface_hub import HfApi, hf_hub_download
from huggingface_hub.utils import disable_progress_bars, enable_progress_bars
disable_progress_bars()
logging.getLogger("huggingface_hub").setLevel(logging.WARNING)
# Find the GGUF file matching the variant
api = HfApi()
files = api.list_repo_files(repo, repo_type="model")
files = api.list_repo_files(repo, repo_type = "model")
gguf_files = [f for f in files if f.endswith(".gguf")]
# Find all GGUF files matching the variant (may be split into shards)
variant_lower = variant.lower().replace("-", "_")
matching = sorted(
f for f in gguf_files
if variant_lower in f.lower().replace("-", "_")
f for f in gguf_files if variant_lower in f.lower().replace("-", "_")
)
if matching:
logger.info(f"Pre-caching helper GGUF: {repo}/{matching[0]}"
+ (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else ""))
logger.info(
f"Pre-caching helper GGUF: {repo}/{matching[0]}"
+ (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else "")
)
for target in matching:
hf_hub_download(repo_id=repo, filename=target)
hf_hub_download(repo_id = repo, filename = target)
logger.info(f"Helper GGUF cached: {len(matching)} file(s)")
else:
logger.warning(f"No GGUF matching variant '{variant}' in {repo}")
except Exception as e:
logger.warning(f"Failed to pre-cache helper GGUF: {e}")
finally:
try:
enable_progress_bars()
except Exception as e:
pass
def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
@ -81,7 +93,9 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
backend = None
try:
@ -89,15 +103,14 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
backend = LlamaCppBackend()
logger.info(f"Loading helper model: {repo} ({variant})")
print(f"🤖 Loading helper model: {repo} ({variant})...")
ok = backend.load_model(
hf_repo=repo,
hf_variant=variant,
model_identifier=f"helper:{repo}:{variant}",
is_vision=False,
n_ctx=2048,
n_gpu_layers=-1,
hf_repo = repo,
hf_variant = variant,
model_identifier = f"helper:{repo}:{variant}",
is_vision = False,
n_ctx = 2048,
n_gpu_layers = -1,
)
if not ok:
logger.warning("Helper model failed to start")
@ -106,12 +119,12 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
messages = [{"role": "user", "content": prompt}]
cumulative = ""
for text in backend.generate_chat_completion(
messages=messages,
temperature=0.1,
top_p=0.9,
top_k=20,
max_tokens=max_tokens,
repetition_penalty=1.0,
messages = messages,
temperature = 0.1,
top_p = 0.9,
top_k = 20,
max_tokens = max_tokens,
repetition_penalty = 1.0,
):
cumulative = text # cumulative — last value is full text
@ -127,7 +140,7 @@ def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
if backend is not None:
try:
backend.unload_model()
print("🤖 Helper model unloaded")
logger.info("Helper model unloaded")
except Exception:
pass
@ -176,7 +189,7 @@ def llm_generate_vlm_instruction(
"Respond with ONLY the instruction sentence, nothing else."
)
result = _run_with_helper(prompt, max_tokens=100)
result = _run_with_helper(prompt, max_tokens = 100)
if not result:
return None
@ -187,7 +200,7 @@ def llm_generate_vlm_instruction(
logger.warning(f"Helper model returned unusable instruction: {instruction!r}")
return None
print(f"🤖 LLM-generated instruction: {instruction}")
logger.info(f"LLM-generated instruction: {instruction}")
return {
"instruction": instruction,
"confidence": 0.85,
@ -231,7 +244,7 @@ def llm_classify_columns(
'Example: {"question": "user", "answer": "assistant", "id": "metadata"}'
)
result = _run_with_helper(prompt, max_tokens=200)
result = _run_with_helper(prompt, max_tokens = 200)
if not result:
return None
@ -248,6 +261,7 @@ def llm_classify_columns(
except json.JSONDecodeError:
# Try to find JSON object in the response
import re
match = re.search(r"\{[^}]+\}", text)
if match:
try:
@ -266,7 +280,11 @@ def llm_classify_columns(
valid_roles = {"user", "assistant", "system", "metadata"}
cleaned = {}
for col, role in mapping.items():
if col in column_names and isinstance(role, str) and role.lower() in valid_roles:
if (
col in column_names
and isinstance(role, str)
and role.lower() in valid_roles
):
cleaned[col] = role.lower()
if not cleaned:
@ -278,7 +296,7 @@ def llm_classify_columns(
logger.warning(f"Helper model mapping missing user/assistant: {cleaned}")
return None
print(f"🤖 LLM-classified columns: {cleaned}")
logger.info(f"LLM-classified columns: {cleaned}")
return cleaned
@ -319,7 +337,7 @@ def llm_generate_dataset_warning(
"Keep it under 3 sentences. Be specific about the dataset."
)
result = _run_with_helper(prompt, max_tokens=200)
result = _run_with_helper(prompt, max_tokens = 200)
if not result:
return None
@ -328,7 +346,7 @@ def llm_generate_dataset_warning(
if len(warning) < 10 or len(warning) > 500:
return None
print(f"🤖 LLM-generated warning: {warning}")
logger.info(f"LLM-generated warning: {warning}")
return warning
@ -369,18 +387,16 @@ def _parse_json_response(text: str) -> Optional[dict]:
return None
def _generate_with_backend(
backend, messages: list[dict], max_tokens: int = 512
) -> str:
def _generate_with_backend(backend, messages: list[dict], max_tokens: int = 512) -> str:
"""Run one chat completion on an already-loaded backend. Returns raw text."""
cumulative = ""
for text in backend.generate_chat_completion(
messages=messages,
temperature=0.1,
top_p=0.9,
top_k=20,
max_tokens=max_tokens,
repetition_penalty=1.0,
messages = messages,
temperature = 0.1,
top_p = 0.9,
top_k = 20,
max_tokens = max_tokens,
repetition_penalty = 1.0,
):
cumulative = text
return cumulative.strip()
@ -398,7 +414,7 @@ def fetch_hf_dataset_card(
try:
from huggingface_hub import DatasetCard
card = DatasetCard.load(dataset_name, token=hf_token)
card = DatasetCard.load(dataset_name, token = hf_token)
readme = card.text or ""
# Truncate at sentence boundary
@ -413,14 +429,21 @@ def fetch_hf_dataset_card(
metadata = {}
if card.data:
for key in (
"task_categories", "task_ids", "language",
"size_categories", "tags", "license", "pretty_name",
"task_categories",
"task_ids",
"language",
"size_categories",
"tags",
"license",
"pretty_name",
):
val = getattr(card.data, key, None)
if val is not None:
metadata[key] = val
logger.info(f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields")
logger.info(
f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields"
)
return readme, metadata
except Exception as e:
@ -447,30 +470,31 @@ def _run_multi_pass_advisor(
return None
repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
variant = os.environ.get(
"UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT
)
backend = None
try:
from core.inference.llama_cpp import LlamaCppBackend
backend = LlamaCppBackend()
print(f"🤖 Loading advisor model: {repo} ({variant})...")
logger.info(f"Loading advisor model: {repo} ({variant})")
t0 = time.monotonic()
ok = backend.load_model(
hf_repo=repo,
hf_variant=variant,
model_identifier=f"advisor:{repo}:{variant}",
is_vision=False,
n_ctx=2048,
n_gpu_layers=-1,
hf_repo = repo,
hf_variant = variant,
model_identifier = f"advisor:{repo}:{variant}",
is_vision = False,
n_ctx = 2048,
n_gpu_layers = -1,
)
if not ok:
logger.warning("Advisor model failed to start")
return None
print(f"🤖 Advisor model loaded in {time.monotonic() - t0:.1f}s")
logger.info(f"Advisor model loaded in {time.monotonic() - t0:.1f}s")
# ── Format samples ──
samples_text = ""
for i, row in enumerate(samples[:5], 1):
@ -478,8 +502,9 @@ def _run_multi_pass_advisor(
samples_text += f"Row {i}:\n" + "\n".join(parts) + "\n"
metadata_str = (
json.dumps(dataset_metadata, indent=2, default=str)[:500]
if dataset_metadata else "N/A"
json.dumps(dataset_metadata, indent = 2, default = str)[:500]
if dataset_metadata
else "N/A"
)
card_excerpt = (dataset_card or "")[:1200] or "N/A"
@ -489,13 +514,14 @@ def _run_multi_pass_advisor(
if model_name:
try:
from utils.models.model_config import load_model_config
config = load_model_config(model_name, use_auth=True, token=hf_token)
config = load_model_config(model_name, use_auth = True, token = hf_token)
archs = getattr(config, "architectures", [])
if archs and "Gemma3nForConditionalGeneration" in archs:
is_gemma_3n = True
except Exception:
is_gemma_3n = "gemma-3n" in model_name.lower()
if model_type == "audio" and not is_gemma_3n:
target_hints = (
"\n\nHINT: The user is training an AUDIO model. The dataset MUST contain "
@ -514,7 +540,7 @@ def _run_multi_pass_advisor(
)
# ── Pass 1: Classify ──
print("🤖 Pass 1: Classifying dataset...", flush=True)
logger.info("Pass 1: Classifying dataset...")
t1 = time.monotonic()
messages1 = [
{
@ -559,9 +585,9 @@ def _run_multi_pass_advisor(
Respond with ONLY the JSON object. No markdown, no explanation."""),
},
]
raw1 = _generate_with_backend(backend, messages1, max_tokens=256)
raw1 = _generate_with_backend(backend, messages1, max_tokens = 256)
pass1 = _parse_json_response(raw1)
print(f"🤖 Pass 1 done ({time.monotonic() - t1:.1f}s): {pass1}", flush=True)
logger.info(f"Pass 1 done ({time.monotonic() - t1:.1f}s): {pass1}")
if not pass1:
logger.warning(f"Advisor Pass 1 failed to produce JSON: {raw1[:200]}")
@ -580,7 +606,9 @@ def _run_multi_pass_advisor(
}
# ── Pass 2: Map columns to roles ──
print("🤖 Pass 2: Mapping columns to roles...", flush=True)
logger.info("Pass 2: Mapping columns to roles...")
t2 = time.monotonic()
messages2 = [
{
@ -592,13 +620,13 @@ def _run_multi_pass_advisor(
'- "user" = This column contains INPUT that the model will receive as a prompt.\n'
'- "assistant" = This column contains OUTPUT that the model should learn to generate.\n\n'
"CRITICAL RULES:\n"
"1. There MUST be at least one column assigned to \"user\" AND at least one "
"column assigned to \"assistant\". Never assign all columns to the same role.\n"
'1. There MUST be at least one column assigned to "user" AND at least one '
'column assigned to "assistant". Never assign all columns to the same role.\n'
"2. The column that contains the TARGET or OUTPUT or ANSWER or LABEL must "
"ALWAYS be assigned to \"assistant\". This is the thing the model should learn "
'ALWAYS be assigned to "assistant". This is the thing the model should learn '
"to produce.\n"
"3. The columns that contain the SOURCE or INPUT or CONTEXT or QUESTION must "
"be assigned to \"user\". This is what the model receives.\n"
'be assigned to "user". This is what the model receives.\n'
'4. Metadata columns like "id", "index", "source", "url", "date" should be '
'set to "skip".\n\n'
"You must respond with ONLY a valid JSON object."
@ -611,7 +639,7 @@ def _run_multi_pass_advisor(
Here is a dataset that has been classified:
CLASSIFICATION:
{json.dumps(pass1, indent=2)}
{json.dumps(pass1, indent = 2)}
COLUMNS AVAILABLE: {columns}
@ -659,9 +687,9 @@ def _run_multi_pass_advisor(
Respond with ONLY the JSON object."""),
},
]
raw2 = _generate_with_backend(backend, messages2, max_tokens=512)
raw2 = _generate_with_backend(backend, messages2, max_tokens = 512)
pass2 = _parse_json_response(raw2)
print(f"🤖 Pass 2 done ({time.monotonic() - t2:.1f}s): {pass2}", flush=True)
logger.info(f"Pass 2 done ({time.monotonic() - t2:.1f}s): {pass2}")
if not pass2:
logger.warning(f"Advisor Pass 2 failed to produce JSON: {raw2[:200]}")
@ -674,10 +702,7 @@ def _run_multi_pass_advisor(
# Validate: must have at least one user AND one assistant
roles_present = set(column_roles.values())
if "user" not in roles_present or "assistant" not in roles_present:
print(
f"🤖 Pass 2 sanity fail: missing user or assistant role: {column_roles}",
flush=True,
)
logger.warning(f"Pass 2 sanity fail: missing user or assistant role: {column_roles}")
return None # triggers fallback to simple classification
# ── Pass 3: System prompt (non-conversational datasets only) ──
@ -686,7 +711,7 @@ def _run_multi_pass_advisor(
is_conv = pass1.get("is_conversational", False)
if not is_conv:
print("🤖 Pass 3: Generating system prompt...", flush=True)
logger.info("Pass 3: Generating system prompt...")
t3 = time.monotonic()
# Format label mapping info for the prompt
@ -726,8 +751,8 @@ def _run_multi_pass_advisor(
Write ONLY the system prompt text. No quotes, no labels, no explanation around it."""),
},
]
raw3 = _generate_with_backend(backend, messages3, max_tokens=256)
print(f"🤖 Pass 3 done ({time.monotonic() - t3:.1f}s): {raw3[:200] if raw3 else None}", flush=True)
raw3 = _generate_with_backend(backend, messages3, max_tokens = 256)
logger.info(f"Pass 3 done ({time.monotonic() - t3:.1f}s): {raw3[:200] if raw3 else None}")
if raw3:
# Pass 3 returns raw text, not JSON — clean it up
@ -746,15 +771,14 @@ def _run_multi_pass_advisor(
note_parts = [f"This is a {dtype} dataset (not conversational)."]
if desc:
note_parts.append(desc)
note_parts.append("Columns have been mapped to conversation roles. You can adjust the mapping if needed.")
note_parts.append(
"Columns have been mapped to conversation roles. You can adjust the mapping if needed."
)
user_notification = " ".join(note_parts)
total_time = time.monotonic() - t0
print(
f"🤖 Advisor complete ({total_time:.1f}s): type={dtype}, "
f"mapping={suggested_mapping}, sys_prompt={bool(sys_prompt)}, label_map={bool(label_map)}",
flush=True,
)
logger.info(f"Advisor complete ({total_time:.1f}s): type={dtype}, mapping={suggested_mapping}, sys_prompt={bool(sys_prompt)}, label_map={bool(label_map)}")
return {
"success": True,
@ -774,7 +798,7 @@ def _run_multi_pass_advisor(
if backend is not None:
try:
backend.unload_model()
print("🤖 Advisor model unloaded")
logger.info("Advisor model unloaded")
except Exception:
pass
@ -805,18 +829,18 @@ def llm_conversion_advisor(
# Try multi-pass advisor
result = _run_multi_pass_advisor(
columns=column_names,
samples=samples,
dataset_name=dataset_name,
dataset_card=dataset_card,
dataset_metadata=dataset_metadata,
model_name=model_name,
model_type=model_type,
hf_token=hf_token,
columns = column_names,
samples = samples,
dataset_name = dataset_name,
dataset_card = dataset_card,
dataset_metadata = dataset_metadata,
model_name = model_name,
model_type = model_type,
hf_token = hf_token,
)
if result and result.get("success"):
print(f"🤖 Conversion advisor succeeded: type={result.get('dataset_type')}")
logger.info(f"Conversion advisor succeeded: type={result.get('dataset_type')}")
return result
# Fallback: simple column classification
@ -826,7 +850,8 @@ def llm_conversion_advisor(
return {
"success": True,
"suggested_mapping": {
col: role for col, role in simple_mapping.items()
col: role
for col, role in simple_mapping.items()
if role in ("user", "assistant", "system")
},
"dataset_type": None,

5
studio/backend/utils/datasets/model_mappings.py
View file

@ -8,7 +8,6 @@ This module contains the mapping dictionaries that associate model names
with their corresponding chat templates and response markers.
"""
TEMPLATE_TO_MODEL_MAPPER = {
"phi-3.5": (
"unsloth/Phi-3.5-mini-instruct-bnb-4bit",
@ -407,14 +406,11 @@ MODEL_TO_TEMPLATE_MAPPER = {}
for key, values in TEMPLATE_TO_MODEL_MAPPER.items():
for value in values:
MODEL_TO_TEMPLATE_MAPPER[value] = key
pass
# Get lowercased
lowered_key = key.lower()
for value in values:
MODEL_TO_TEMPLATE_MAPPER[value.lower()] = lowered_key
pass
pass
TEMPLATE_TO_RESPONSES_MAPPER = {
@ -531,4 +527,3 @@ TEMPLATE_TO_RESPONSES_MAPPER = {
"response": "<|assistant|><think>",
},
}

41
studio/backend/utils/datasets/vlm_processing.py
View file

@ -14,9 +14,9 @@ from itertools import islice
def generate_smart_vlm_instruction(
dataset,
text_column="text",
image_column="image",
dataset_name=None,
text_column = "text",
image_column = "image",
dataset_name = None,
):
"""
Generate smart, context-aware instruction for VLM datasets using heuristics.
@ -66,11 +66,12 @@ def generate_smart_vlm_instruction(
# OCR / Transcription
"ocr": {
"keywords": ["ocr", "transcribe", "transcript"],
"content_hints": [r"[A-Za-z\u0600-\u06FF]{10,}"], # Long text passages (Latin/Arabic)
"content_hints": [
r"[A-Za-z\u0600-\u06FF]{10,}"
], # Long text passages (Latin/Arabic)
"instruction": "Transcribe all the text shown in this image.",
"confidence": 0.9,
},
# LaTeX / Math
"latex": {
"keywords": ["latex", "math", "formula", "equation"],
@ -78,7 +79,6 @@ def generate_smart_vlm_instruction(
"instruction": "Convert this image to LaTeX notation.",
"confidence": 0.95,
},
# Caption / Description
"caption": {
"keywords": ["caption", "description", "describe"],
@ -86,15 +86,21 @@ def generate_smart_vlm_instruction(
"instruction": "Provide a detailed description of this image.",
"confidence": 0.85,
},
# Medical / Radiology
"medical": {
"keywords": ["medical", "radiology", "xray", "ct", "mri", "scan", "diagnosis"],
"keywords": [
"medical",
"radiology",
"xray",
"ct",
"mri",
"scan",
"diagnosis",
],
"content_hints": [r"\b(lesion|radiograph|patient|diagnosis|findings)\b"],
"instruction": "Analyze this medical image and describe the key findings.",
"confidence": 0.9,
},
# Code / Programming
"code": {
"keywords": ["code", "program", "function", "algorithm"],
@ -102,7 +108,6 @@ def generate_smart_vlm_instruction(
"instruction": "Explain what this code visualization shows.",
"confidence": 0.85,
},
# Chart / Graph
"chart": {
"keywords": ["chart", "graph", "plot", "visualization", "diagram"],
@ -110,7 +115,6 @@ def generate_smart_vlm_instruction(
"instruction": "Describe this chart or graph, including key data points and trends.",
"confidence": 0.85,
},
# Document / Text Recognition
"document": {
"keywords": ["document", "page", "paragraph", "article"],
@ -132,7 +136,9 @@ def generate_smart_vlm_instruction(
score += 0.5
# Check dataset name if provided
if dataset_name and any(keyword in dataset_name.lower() for keyword in task_info["keywords"]):
if dataset_name and any(
keyword in dataset_name.lower() for keyword in task_info["keywords"]
):
score += 0.3
# Check content patterns
@ -186,7 +192,7 @@ def generate_smart_vlm_instruction(
row = {}
for col in s:
val = s[col]
if hasattr(val, 'size') and hasattr(val, 'mode'): # PIL Image
if hasattr(val, "size") and hasattr(val, "mode"): # PIL Image
row[col] = "<image>"
elif isinstance(val, list):
row[col] = str(val)[:300]
@ -195,15 +201,15 @@ def generate_smart_vlm_instruction(
sample_rows.append(row)
llm_result = llm_generate_vlm_instruction(
column_names=list(column_names),
samples=sample_rows,
dataset_name=dataset_name,
column_names = list(column_names),
samples = sample_rows,
dataset_name = dataset_name,
)
if llm_result and llm_result.get("instruction"):
print(
f"\n[DEBUG] LLM-assisted VLM instruction generated: "
f"'{llm_result['instruction']}' (confidence={llm_result.get('confidence', 'N/A')})\n",
flush=True,
flush = True,
)
return {
"instruction": llm_result["instruction"],
@ -214,6 +220,7 @@ def generate_smart_vlm_instruction(
}
except Exception as e:
import logging
logging.getLogger(__name__).debug(f"LLM-assisted instruction skipped: {e}")
# ===== LEVEL 5: Generic Fallback =====

27
studio/backend/utils/hardware/__init__.py
View file

@ -4,6 +4,7 @@
"""
Hardware detection and GPU utilities
"""
from .hardware import (
DeviceType,
DEVICE,
@ -21,17 +22,17 @@ from .hardware import (
)
__all__ = [
'DeviceType',
'DEVICE',
'detect_hardware',
'get_device',
'is_apple_silicon',
'clear_gpu_cache',
'get_gpu_memory_info',
'log_gpu_memory',
'get_gpu_summary',
'get_package_versions',
'get_gpu_utilization',
'get_physical_gpu_count',
'safe_num_proc',
"DeviceType",
"DEVICE",
"detect_hardware",
"get_device",
"is_apple_silicon",
"clear_gpu_cache",
"get_gpu_memory_info",
"log_gpu_memory",
"get_gpu_summary",
"get_package_versions",
"get_gpu_utilization",
"get_physical_gpu_count",
"safe_num_proc",
]

34
studio/backend/utils/hardware/hardware.py
View file

@ -15,6 +15,7 @@ Usage:
import torch
...
"""
import platform
import structlog
from loggers import get_logger
@ -26,11 +27,13 @@ logger = get_logger(__name__)
# ========== Device Enum ==========
class DeviceType(str, Enum):
"""Supported compute backends. Inherits from str so it serializes cleanly in JSON."""
CUDA = "cuda"
MLX = "mlx"
CPU = "cpu"
MLX = "mlx"
CPU = "cpu"
# ========== Global State (set once by detect_hardware) ==========
@ -40,6 +43,7 @@ DEVICE: Optional[DeviceType] = None
# ========== Detection ==========
def is_apple_silicon() -> bool:
"""Check if running on Apple Silicon hardware (pure platform check, no ML imports)."""
return platform.system() == "Darwin" and platform.machine() == "arm64"
@ -49,6 +53,7 @@ def _has_torch() -> bool:
"""Check if PyTorch is importable."""
try:
import torch
return True
except ImportError:
return False
@ -58,6 +63,7 @@ def _has_mlx() -> bool:
"""Check if MLX is importable."""
try:
import mlx.core
return True
except ImportError:
return False
@ -80,6 +86,7 @@ def detect_hardware() -> DeviceType:
# --- CUDA: try PyTorch ---
if _has_torch():
import torch
if torch.cuda.is_available():
DEVICE = DeviceType.CUDA
device_name = torch.cuda.get_device_properties(0).name
@ -101,6 +108,7 @@ def detect_hardware() -> DeviceType:
# ========== Convenience helpers ==========
def get_device() -> DeviceType:
"""
Return the detected device. Auto-detects if detect_hardware() hasn't been called yet.
@ -118,12 +126,14 @@ def clear_gpu_cache():
Safe to call on any platform — no-ops gracefully.
"""
import gc
gc.collect()
device = get_device()
if device == DeviceType.CUDA:
import torch
torch.cuda.synchronize()
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
@ -144,6 +154,7 @@ def get_gpu_memory_info() -> Dict[str, Any]:
if device == DeviceType.CUDA:
try:
import torch
idx = torch.cuda.current_device()
props = torch.cuda.get_device_properties(idx)
@ -215,6 +226,7 @@ def log_gpu_memory(context: str):
# ========== GPU Summary & Package Versions ==========
def get_gpu_summary() -> Dict[str, Any]:
"""
Return a compact summary of the primary GPU.
@ -256,6 +268,7 @@ def get_package_versions() -> Dict[str, Optional[str]]:
# CUDA toolkit version bundled with torch
try:
import torch
versions["cuda"] = getattr(torch.version, "cuda", None)
except Exception:
versions["cuda"] = None
@ -265,6 +278,7 @@ def get_package_versions() -> Dict[str, Optional[str]]:
# ========== Live GPU Utilization (nvidia-smi) ==========
def get_gpu_utilization() -> Dict[str, Any]:
"""
Return a live snapshot of GPU utilization via ``nvidia-smi``.
@ -312,9 +326,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
"memory.used,memory.total,power.draw,power.limit",
"--format=csv,noheader,nounits",
],
capture_output=True,
text=True,
timeout=5,
capture_output = True,
text = True,
timeout = 5,
)
if result.returncode == 0 and result.stdout.strip():
@ -360,7 +374,9 @@ def get_gpu_utilization() -> Dict[str, Any]:
power_limit = smi_data.get("power_limit")
vram_used_gb = round(vram_used_mb / 1024, 2) if vram_used_mb is not None else None
vram_total_gb = round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
vram_total_gb = (
round(vram_total_mb / 1024, 2) if vram_total_mb is not None else None
)
vram_pct = (
round((vram_used_mb / vram_total_mb) * 100, 1)
if vram_used_mb is not None and vram_total_mb and vram_total_mb > 0
@ -395,6 +411,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
_physical_gpu_count: Optional[int] = None
def get_physical_gpu_count() -> int:
"""
Return the number of physical NVIDIA GPUs on the machine.
@ -409,9 +426,12 @@ def get_physical_gpu_count() -> int:
try:
import subprocess
result = subprocess.run(
["nvidia-smi", "-L"],
capture_output=True, text=True, timeout=5,
capture_output = True,
text = True,
timeout = 5,
)
if result.returncode == 0 and result.stdout.strip():
_physical_gpu_count = len(result.stdout.strip().splitlines())

2
studio/backend/utils/inference/__init__.py
View file

@ -4,7 +4,7 @@
"""
Inference utility functions
"""
from utils.inference.inference_config import load_inference_config
__all__ = ["load_inference_config"]

26
studio/backend/utils/inference/inference_config.py
View file

@ -7,6 +7,7 @@ Inference configuration loading utilities.
This module provides functions to load inference parameters (temperature, top_p, top_k, min_p)
from model YAML configuration files, with fallback to default.yaml.
"""
from pathlib import Path
from typing import Dict, Any
import yaml
@ -21,15 +22,15 @@ logger = get_logger(__name__)
def load_inference_config(model_identifier: str) -> Dict[str, Any]:
"""
Load inference configuration parameters for a model.
This function loads inference parameters (temperature, top_p, top_k, min_p) from the
model's YAML configuration file using the same mapping logic as the /config endpoint.
If a parameter is missing from the model's config, it falls back to the value in
default.yaml.
Args:
model_identifier: Model identifier (e.g., "unsloth/llama-3-8b-bnb-4bit")
Returns:
Dictionary containing inference parameters:
{
@ -41,30 +42,33 @@ def load_inference_config(model_identifier: str) -> Dict[str, Any]:
"""
# Load model defaults to get inference parameters
model_defaults = load_model_defaults(model_identifier)
# Load default.yaml for fallback values
script_dir = Path(__file__).parent.parent.parent
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
default_config_path = defaults_dir / "default.yaml"
default_inference = {}
if default_config_path.exists():
try:
with open(default_config_path, 'r', encoding='utf-8') as f:
with open(default_config_path, "r", encoding = "utf-8") as f:
default_config = yaml.safe_load(f) or {}
default_inference = default_config.get("inference", {})
except Exception as e:
logger.warning(f"Failed to load default.yaml: {e}")
# Extract inference parameters from model config, fallback to defaults
model_inference = model_defaults.get("inference", {})
inference_config = {
"temperature": model_inference.get("temperature", default_inference.get("temperature", 0.7)),
"temperature": model_inference.get(
"temperature", default_inference.get("temperature", 0.7)
),
"top_p": model_inference.get("top_p", default_inference.get("top_p", 0.95)),
"top_k": model_inference.get("top_k", default_inference.get("top_k", -1)),
"min_p": model_inference.get("min_p", default_inference.get("min_p", 0.01)),
"trust_remote_code": model_inference.get("trust_remote_code", default_inference.get("trust_remote_code", False)),
"trust_remote_code": model_inference.get(
"trust_remote_code", default_inference.get("trust_remote_code", False)
),
}
return inference_config
return inference_config

33
studio/backend/utils/models/__init__.py
View file

@ -4,6 +4,7 @@
"""
Model and LoRA configuration handling
"""
from .model_config import (
ModelConfig,
GgufVariantInfo,
@ -24,20 +25,20 @@ from .model_config import (
from .checkpoints import scan_checkpoints
__all__ = [
'ModelConfig',
'GgufVariantInfo',
'is_vision_model',
'is_embedding_model',
'detect_audio_type',
'is_audio_input_type',
'VALID_AUDIO_TYPES',
'scan_trained_loras',
'scan_exported_models',
'load_model_defaults',
'get_base_model_from_lora',
'load_model_config',
'list_gguf_variants',
'MODEL_NAME_MAPPING',
'UI_STATUS_INDICATORS',
'scan_checkpoints',
"ModelConfig",
"GgufVariantInfo",
"is_vision_model",
"is_embedding_model",
"detect_audio_type",
"is_audio_input_type",
"VALID_AUDIO_TYPES",
"scan_trained_loras",
"scan_exported_models",
"load_model_defaults",
"get_base_model_from_lora",
"load_model_config",
"list_gguf_variants",
"MODEL_NAME_MAPPING",
"UI_STATUS_INDICATORS",
"scan_checkpoints",
]

17
studio/backend/utils/models/checkpoints.py
View file

@ -4,6 +4,7 @@
"""
Checkpoint scanning utilities for discovering training runs and their checkpoints.
"""
import json
import structlog
from loggers import get_logger
@ -86,7 +87,9 @@ def scan_checkpoints(
name_part = parts[0]
idx = name_part.find("_")
if idx > 0:
metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1:]
metadata["base_model"] = (
name_part[:idx] + "/" + name_part[idx + 1 :]
)
else:
metadata["base_model"] = name_part
@ -109,13 +112,19 @@ def scan_checkpoints(
# Assign the last checkpoint's loss to the main adapter entry
if len(checkpoints) > 1:
last_checkpoint_loss = checkpoints[-1][2]
checkpoints[0] = (checkpoints[0][0], checkpoints[0][1], last_checkpoint_loss)
checkpoints[0] = (
checkpoints[0][0],
checkpoints[0][1],
last_checkpoint_loss,
)
models.append((item.name, checkpoints, metadata))
logger.debug(f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)")
logger.debug(
f"Found model: {item.name} with {len(checkpoints)} checkpoint(s)"
)
# Sort by modification time (newest first)
models.sort(key=lambda x: Path(x[1][0][1]).stat().st_mtime, reverse=True)
models.sort(key = lambda x: Path(x[1][0][1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(models)} training runs in {outputs_dir}")
return models

512
studio/backend/utils/models/model_config.py
View file

@ -4,6 +4,7 @@
"""
Model and LoRA configuration handling
"""
from transformers import AutoConfig
from dataclasses import dataclass
from typing import Optional, Dict, Any
@ -77,7 +78,6 @@ MODEL_NAME_MAPPING = {
"unsloth_ERNIE-4.5-VL-28B-A3B-PT.yaml": [
"unsloth/ERNIE-4.5-VL-28B-A3B-PT",
],
"tiiuae_Falcon-H1-0.5B-Instruct.yaml": [
"tiiuae/Falcon-H1-0.5B-Instruct",
"unsloth/Falcon-H1-0.5B-Instruct",
@ -132,7 +132,6 @@ MODEL_NAME_MAPPING = {
"unsloth/gpt-oss-20b-unsloth-bnb-4bit",
"unsloth/gpt-oss-20b-BF16",
],
"unsloth_gpt-oss-120b.yaml": [
"openai/gpt-oss-120b",
"unsloth/gpt-oss-120b-unsloth-bnb-4bit",
@ -169,7 +168,6 @@ MODEL_NAME_MAPPING = {
"unsloth/Meta-Llama-3.1-405B-bnb-4bit",
"meta-llama/Meta-Llama-3.1-405B",
],
"unsloth_Meta-Llama-3.1-8B-Instruct-bnb-4bit.yaml": [
"unsloth/Meta-Llama-3.1-8B-Instruct-unsloth-bnb-4bit",
"unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
@ -229,10 +227,9 @@ MODEL_NAME_MAPPING = {
"unsloth/Mistral-Nemo-Base-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Base-2407",
"mistralai/Mistral-Nemo-Base-2407",
"unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Instruct-2407",
"unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
"unsloth/Mistral-Nemo-Instruct-2407",
"mistralai/Mistral-Nemo-Instruct-2407",
],
"unsloth_Mistral-Small-Instruct-2409.yaml": [
"unsloth/Mistral-Small-Instruct-2409-bnb-4bit",
@ -384,7 +381,10 @@ for canonical_file, model_names in MODEL_NAME_MAPPING.items():
for model_name in model_names:
_REVERSE_MODEL_MAPPING[model_name.lower()] = canonical_file
def load_model_config(model_name: str, use_auth: bool = False, token: Optional[str] = None):
def load_model_config(
model_name: str, use_auth: bool = False, token: Optional[str] = None
):
"""
Load model config with optional authentication control.
"""
@ -392,32 +392,30 @@ def load_model_config(model_name: str, use_auth: bool = False, token: Optional[s
if token:
# Explicit token provided - use it
return AutoConfig.from_pretrained(
model_name,
trust_remote_code=True,
token=token
model_name, trust_remote_code = True, token = token
)
if not use_auth:
# Load without any authentication (for public model checks)
with without_hf_auth():
return AutoConfig.from_pretrained(
model_name,
trust_remote_code=True,
token=None
model_name, trust_remote_code = True, token = None
)
# Use default authentication (cached tokens)
return AutoConfig.from_pretrained(
model_name,
trust_remote_code=True
)
return AutoConfig.from_pretrained(model_name, trust_remote_code = True)
# VLM architecture suffixes and known VLM model_type values.
_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text")
_VLM_MODEL_TYPES = {
'phi3_v', 'llava', 'llava_next', 'llava_onevision',
'internvl_chat', 'cogvlm2', 'minicpmv',
"phi3_v",
"llava",
"llava_next",
"llava_onevision",
"internvl_chat",
"cogvlm2",
"minicpmv",
}
# Pre-computed .venv_t5 path and backend dir for subprocess version switching.
@ -426,7 +424,7 @@ _BACKEND_DIR = str(Path(__file__).resolve().parent.parent.parent)
# Inline script executed in a subprocess with transformers 5.x activated.
# Receives model_name and token via argv, prints JSON result to stdout.
_VISION_CHECK_SCRIPT = r'''
_VISION_CHECK_SCRIPT = r"""
import sys, os, json
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@ -472,10 +470,12 @@ try:
except Exception as exc:
logger.info(json.dumps({"error": str(exc)}))
sys.exit(1)
'''
"""
def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> bool:
def _is_vision_model_subprocess(
model_name: str, hf_token: Optional[str] = None
) -> bool:
"""Run is_vision_model check in a subprocess with transformers 5.x.
Same pattern as training/inference workers: spawn a clean subprocess
@ -486,16 +486,26 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
try:
result = subprocess.run(
[sys.executable, "-c", _VISION_CHECK_SCRIPT,
_VENV_T5_DIR, _BACKEND_DIR, model_name, token_arg],
capture_output=True, text=True, timeout=60,
[
sys.executable,
"-c",
_VISION_CHECK_SCRIPT,
_VENV_T5_DIR,
_BACKEND_DIR,
model_name,
token_arg,
],
capture_output = True,
text = True,
timeout = 60,
)
if result.returncode != 0:
stderr = result.stderr.strip()
logger.warning(
"Vision check subprocess failed for '%s': %s",
model_name, stderr or result.stdout.strip(),
model_name,
stderr or result.stdout.strip(),
)
return False
@ -503,7 +513,8 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
if "error" in data:
logger.warning(
"Vision check subprocess error for '%s': %s",
model_name, data["error"],
model_name,
data["error"],
)
return False
@ -511,7 +522,10 @@ def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None)
logger.info(
"Vision check (subprocess, transformers 5.x) for '%s': "
"model_type=%s, architectures=%s, is_vision=%s",
model_name, data.get("model_type"), data.get("architectures"), is_vlm,
model_name,
data.get("model_type"),
data.get("architectures"),
is_vlm,
)
return is_vlm
@ -540,52 +554,54 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
# because AutoConfig in the main process (transformers 4.57.x) doesn't
# recognize their architectures.
from utils.transformers_version import needs_transformers_5
if needs_transformers_5(model_name):
logger.info(
"Model '%s' needs transformers 5.x — checking vision via subprocess",
model_name,
)
return _is_vision_model_subprocess(model_name, hf_token=hf_token)
return _is_vision_model_subprocess(model_name, hf_token = hf_token)
try:
config = load_model_config(model_name, use_auth=True, token=hf_token)
config = load_model_config(model_name, use_auth = True, token = hf_token)
# Exclude audio-only models that share ForConditionalGeneration suffix
# (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
_audio_only_model_types = {'csm', 'whisper'}
model_type = getattr(config, 'model_type', None)
_audio_only_model_types = {"csm", "whisper"}
model_type = getattr(config, "model_type", None)
if model_type in _audio_only_model_types:
return False
# Check 1: Architecture class name patterns
if hasattr(config, 'architectures'):
is_vlm = any(
x.endswith(_VLM_ARCH_SUFFIXES)
for x in config.architectures
)
if hasattr(config, "architectures"):
is_vlm = any(x.endswith(_VLM_ARCH_SUFFIXES) for x in config.architectures)
if is_vlm:
logger.info(f"Model {model_name} detected as VLM: architecture {config.architectures}")
logger.info(
f"Model {model_name} detected as VLM: architecture {config.architectures}"
)
return True
# Check 2: Has vision_config (most VLMs: LLaVA, Gemma-3, Qwen2-VL, etc.)
if hasattr(config, 'vision_config'):
if hasattr(config, "vision_config"):
logger.info(f"Model {model_name} detected as VLM: has vision_config")
return True
# Check 3: Has img_processor (Phi-3.5 Vision uses this instead of vision_config)
if hasattr(config, 'img_processor'):
if hasattr(config, "img_processor"):
logger.info(f"Model {model_name} detected as VLM: has img_processor")
return True
# Check 4: Has image_token_index (common in VLMs for image placeholder tokens)
if hasattr(config, 'image_token_index'):
if hasattr(config, "image_token_index"):
logger.info(f"Model {model_name} detected as VLM: has image_token_index")
return True
# Check 5: Known VLM model_type values that may not match above checks
if hasattr(config, 'model_type'):
if hasattr(config, "model_type"):
if config.model_type in _VLM_MODEL_TYPES:
logger.info(f"Model {model_name} detected as VLM: model_type={config.model_type}")
logger.info(
f"Model {model_name} detected as VLM: model_type={config.model_type}"
)
return True
return False
@ -595,19 +611,20 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
return False
VALID_AUDIO_TYPES = ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm')
VALID_AUDIO_TYPES = ("snac", "csm", "bicodec", "dac", "whisper", "audio_vlm")
# Cache detection results per session to avoid repeated API calls
_audio_detection_cache: Dict[str, Optional[str]] = {}
# Tokenizer token patterns → audio_type (all 6 types detected from tokenizer_config.json)
_AUDIO_TOKEN_PATTERNS = {
'csm': lambda tokens: '<|AUDIO|>' in tokens and '<|audio_eos|>' in tokens,
'whisper': lambda tokens: '<|startoftranscript|>' in tokens,
'audio_vlm': lambda tokens: '<audio_soft_token>' in tokens,
'bicodec': lambda tokens: any(t.startswith('<|bicodec_') for t in tokens),
'dac': lambda tokens: '<|audio_start|>' in tokens and '<|audio_end|>' in tokens,
'snac': lambda tokens: sum(1 for t in tokens if t.startswith('<custom_token_')) > 10000,
"csm": lambda tokens: "<|AUDIO|>" in tokens and "<|audio_eos|>" in tokens,
"whisper": lambda tokens: "<|startoftranscript|>" in tokens,
"audio_vlm": lambda tokens: "<audio_soft_token>" in tokens,
"bicodec": lambda tokens: any(t.startswith("<|bicodec_") for t in tokens),
"dac": lambda tokens: "<|audio_start|>" in tokens and "<|audio_end|>" in tokens,
"snac": lambda tokens: sum(1 for t in tokens if t.startswith("<custom_token_"))
> 10000,
}
@ -631,17 +648,20 @@ def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Option
return result
def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
def _detect_audio_from_tokenizer(
model_name: str, hf_token: Optional[str] = None
) -> Optional[str]:
"""Detect audio type from tokenizer special tokens (for LLM-based audio models).
First checks local HF cache, then fetches tokenizer_config.json from HuggingFace.
Checks added_tokens_decoder for distinctive patterns.
"""
def _check_token_patterns(tok_config: dict) -> Optional[str]:
added = tok_config.get('added_tokens_decoder', {})
added = tok_config.get("added_tokens_decoder", {})
if not added:
return None
token_contents = [v.get('content', '') for v in added.values()]
token_contents = [v.get("content", "") for v in added.values()]
for audio_type, check_fn in _AUDIO_TOKEN_PATTERNS.items():
if check_fn(token_contents):
return audio_type
@ -650,6 +670,7 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
# 1) Check local HF cache first (works for gated/offline models)
try:
from huggingface_hub.constants import HF_HUB_CACHE
cache_dir = Path(HF_HUB_CACHE)
repo_dir_name = f"models--{model_name.replace('/', '--')}"
repo_dir = cache_dir / repo_dir_name
@ -657,7 +678,10 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
snapshots_dir = repo_dir / "snapshots"
if snapshots_dir.exists():
for snapshot in snapshots_dir.iterdir():
for tok_path in ['tokenizer_config.json', 'LLM/tokenizer_config.json']:
for tok_path in [
"tokenizer_config.json",
"LLM/tokenizer_config.json",
]:
tok_file = snapshot / tok_path
if tok_file.exists():
tok_config = json.loads(tok_file.read_text())
@ -672,16 +696,16 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
import requests
import os
paths_to_try = ['tokenizer_config.json', 'LLM/tokenizer_config.json']
paths_to_try = ["tokenizer_config.json", "LLM/tokenizer_config.json"]
# Use provided token, or fall back to env
token = hf_token or os.environ.get('HF_TOKEN')
token = hf_token or os.environ.get("HF_TOKEN")
headers = {}
if token:
headers['Authorization'] = f'Bearer {token}'
headers["Authorization"] = f"Bearer {token}"
for tok_path in paths_to_try:
url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
resp = requests.get(url, headers=headers, timeout=15)
resp = requests.get(url, headers = headers, timeout = 15)
if not resp.ok:
continue
@ -692,7 +716,9 @@ def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None
return None
except Exception as e:
logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}")
logger.debug(
f"Could not detect audio type from tokenizer for {model_name}: {e}"
)
return None
@ -701,7 +727,7 @@ def is_audio_input_type(audio_type: Optional[str]) -> bool:
Whisper (ASR) and audio_vlm (Gemma3n) accept audio input.
"""
return audio_type in ('whisper', 'audio_vlm')
return audio_type in ("whisper", "audio_vlm")
def _is_mmproj(filename: str) -> bool:
@ -756,7 +782,8 @@ def detect_gguf_model(path: str) -> Optional[str]:
if p.is_dir():
gguf_files = sorted(
(f for f in p.glob("*.gguf") if not _is_mmproj(f.name)),
key=lambda f: f.stat().st_size, reverse=True,
key = lambda f: f.stat().st_size,
reverse = True,
)
if gguf_files:
return str(gguf_files[0].resolve())
@ -767,9 +794,18 @@ def detect_gguf_model(path: str) -> Optional[str]:
# Preferred GGUF quantization levels, in descending priority.
# Q4_K_M is a good default: small, fast, acceptable quality.
_GGUF_QUANT_PREFERENCE = [
"Q4_K_M", "Q4_K_S", "Q5_K_M", "Q5_K_S",
"Q6_K", "Q8_0", "Q3_K_M", "Q3_K_L", "Q2_K",
"F16", "BF16", "F32",
"Q4_K_M",
"Q4_K_S",
"Q5_K_M",
"Q5_K_S",
"Q6_K",
"Q8_0",
"Q3_K_M",
"Q3_K_L",
"Q2_K",
"F16",
"BF16",
"F32",
]
@ -797,9 +833,10 @@ def _pick_best_gguf(filenames: list[str]) -> Optional[str]:
@dataclass
class GgufVariantInfo:
"""A single GGUF quantization variant from a HuggingFace repo."""
filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
quant: str # e.g., "Q4_K_M" (extracted from filename)
size_bytes: int # file size
filename: str # e.g., "gemma-3-4b-it-Q4_K_M.gguf"
quant: str # e.g., "Q4_K_M" (extracted from filename)
size_bytes: int # file size
def _extract_quant_label(filename: str) -> str:
@ -815,21 +852,23 @@ def _extract_quant_label(filename: str) -> str:
"MXFP4_MOE/model-MXFP4_MOE-0001.gguf"→ "MXFP4_MOE"
"""
import re
# Use only the basename (rfilename may include directory)
basename = filename.rsplit("/", 1)[-1]
# Strip .gguf and any shard suffix (-00001-of-00010)
stem = re.sub(r'-\d{3,}-of-\d{3,}', '', basename.rsplit(".", 1)[0])
stem = re.sub(r"-\d{3,}-of-\d{3,}", "", basename.rsplit(".", 1)[0])
# Match known quantization patterns
match = re.search(
r'(UD-)?' # Optional UD- prefix (Ultra Discrete)
r'(MXFP[0-9]+(?:_[A-Z0-9]+)*' # MXFP variants: MXFP4, MXFP4_MOE
r'|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?' # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r'|TQ[0-9]+_[0-9]+' # Ternary quant: TQ1_0, TQ2_0
r'|Q[0-9]+_K_[A-Z]+' # K-quant: Q4_K_M, Q3_K_S
r'|Q[0-9]+_[0-9]+' # Standard: Q8_0, Q5_1
r'|Q[0-9]+_K' # Short K-quant: Q6_K
r'|BF16|F16|F32)', # Full precision
stem, re.IGNORECASE,
r"(UD-)?" # Optional UD- prefix (Ultra Discrete)
r"(MXFP[0-9]+(?:_[A-Z0-9]+)*" # MXFP variants: MXFP4, MXFP4_MOE
r"|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?" # IQ variants: IQ4_XS, IQ4_NL, IQ1_S
r"|TQ[0-9]+_[0-9]+" # Ternary quant: TQ1_0, TQ2_0
r"|Q[0-9]+_K_[A-Z]+" # K-quant: Q4_K_M, Q3_K_S
r"|Q[0-9]+_[0-9]+" # Standard: Q8_0, Q5_1
r"|Q[0-9]+_K" # Short K-quant: Q6_K
r"|BF16|F16|F32)", # Full precision
stem,
re.IGNORECASE,
)
if match:
prefix = match.group(1) or ""
@ -853,11 +892,11 @@ def list_gguf_variants(
"""
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(repo_id, token=hf_token, files_metadata=True)
info = hf_model_info(repo_id, token = hf_token, files_metadata = True)
variants: list[GgufVariantInfo] = []
has_vision = False
quant_totals: dict[str, int] = {} # quant -> total bytes
quant_totals: dict[str, int] = {} # quant -> total bytes
quant_first_file: dict[str, str] = {} # quant -> first filename (for display)
for sibling in info.siblings:
@ -877,11 +916,13 @@ def list_gguf_variants(
quant_first_file[quant] = fname
for quant, total_size in quant_totals.items():
variants.append(GgufVariantInfo(
filename=quant_first_file[quant],
quant=quant,
size_bytes=total_size,
))
variants.append(
GgufVariantInfo(
filename = quant_first_file[quant],
quant = quant,
size_bytes = total_size,
)
)
return variants, has_vision
@ -898,7 +939,7 @@ def detect_gguf_model_remote(
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(repo_id, token=hf_token)
info = hf_model_info(repo_id, token = hf_token)
repo_files = [s.rfilename for s in info.siblings]
return _pick_best_gguf(repo_files)
except Exception as e:
@ -919,9 +960,9 @@ def download_gguf_file(
from huggingface_hub import hf_hub_download
local_path = hf_hub_download(
repo_id=repo_id,
filename=filename,
token=hf_token,
repo_id = repo_id,
filename = filename,
token = hf_token,
)
return local_path
@ -964,7 +1005,7 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool:
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(model_name, token=hf_token)
info = hf_model_info(model_name, token = hf_token)
tags = set(info.tags or [])
pipeline_tag = info.pipeline_tag or ""
@ -1024,16 +1065,21 @@ def scan_trained_loras(outputs_dir: str = str(outputs_root())) -> List[Tuple[str
logger.debug(f"Found trained LoRA: {display_name}")
# Sort by modification time (newest first)
trained_loras.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
trained_loras.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}")
logger.info(
f"Found {len(trained_loras)} trained LoRA adapters in {outputs_dir}"
)
return trained_loras
except Exception as e:
logger.error(f"Error scanning outputs folder: {e}")
return []
def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[str, str, str, Optional[str]]]:
def scan_exported_models(
exports_dir: str = str(exports_root()),
) -> List[Tuple[str, str, str, Optional[str]]]:
"""
Scan exports folder for exported models (merged, LoRA, GGUF).
@ -1082,9 +1128,8 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
adapter_config = checkpoint_dir / "adapter_config.json"
config_file = checkpoint_dir / "config.json"
has_weights = (
any(checkpoint_dir.glob("*.safetensors"))
or any(checkpoint_dir.glob("*.bin"))
has_weights = any(checkpoint_dir.glob("*.safetensors")) or any(
checkpoint_dir.glob("*.bin")
)
has_gguf = any(checkpoint_dir.glob("*.gguf"))
@ -1134,7 +1179,9 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
# Fallback: read base model from the original training run's
# adapter_config.json in ./outputs/{run_name}/
if not base_model:
outputs_adapter_cfg = resolve_output_dir(run_dir.name) / "adapter_config.json"
outputs_adapter_cfg = (
resolve_output_dir(run_dir.name) / "adapter_config.json"
)
try:
if outputs_adapter_cfg.exists():
cfg = json.loads(outputs_adapter_cfg.read_text())
@ -1147,7 +1194,7 @@ def scan_exported_models(exports_dir: str = str(exports_root())) -> List[Tuple[s
results.append((display_name, model_path, export_type, base_model))
logger.debug(f"Found exported model: {display_name} ({export_type})")
results.sort(key=lambda x: Path(x[1]).stat().st_mtime, reverse=True)
results.sort(key = lambda x: Path(x[1]).stat().st_mtime, reverse = True)
logger.info(f"Found {len(results)} exported models in {exports_dir}")
return results
@ -1177,11 +1224,13 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
# Try adapter_config.json first
adapter_config_path = lora_path_obj / "adapter_config.json"
if adapter_config_path.exists():
with open(adapter_config_path, 'r') as f:
with open(adapter_config_path, "r") as f:
config = json.load(f)
base_model = config.get("base_model_name_or_path")
if base_model:
logger.info(f"Detected base model from adapter_config.json: {base_model}")
logger.info(
f"Detected base model from adapter_config.json: {base_model}"
)
return base_model
# Fallback: try training_args.bin (requires torch)
@ -1189,10 +1238,13 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
if training_args_path.exists():
try:
import torch
training_args = torch.load(training_args_path)
if hasattr(training_args, 'model_name_or_path'):
if hasattr(training_args, "model_name_or_path"):
base_model = training_args.model_name_or_path
logger.info(f"Detected base model from training_args.bin: {base_model}")
logger.info(
f"Detected base model from training_args.bin: {base_model}"
)
return base_model
except Exception as e:
logger.warning(f"Could not load training_args.bin: {e}")
@ -1216,22 +1268,23 @@ def get_base_model_from_lora(lora_path: str) -> Optional[str]:
except Exception as e:
logger.error(f"Error reading base model from LoRA config: {e}")
return None
pass
# Status indicators that appear in UI dropdowns
UI_STATUS_INDICATORS = [" (Ready)", " (Loading...)", " (Active)", "↓ "]
def load_model_defaults(model_name: str) -> Dict[str, Any]:
"""
Load default training parameters for a model from YAML file.
Args:
model_name: Model identifier (e.g., "unsloth/Meta-Llama-3.1-8B-bnb-4bit")
Returns:
Dictionary with default parameters from YAML file, or empty dict if not found
The function looks for a YAML file in configs/model_defaults/ (including subfolders)
The function looks for a YAML file in configs/model_defaults/ (including subfolders)
based on the model name or its aliases from MODEL_NAME_MAPPING.
If no specific file exists, it falls back to default.yaml.
"""
@ -1239,22 +1292,26 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
# Get the script directory to locate configs
script_dir = Path(__file__).parent.parent.parent
defaults_dir = script_dir / "assets" / "configs" / "model_defaults"
# First, check if model is in the mapping
if model_name.lower() in _REVERSE_MODEL_MAPPING:
canonical_file = _REVERSE_MODEL_MAPPING[model_name.lower()]
# Search in subfolders and root
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path} (via mapping)")
logger.info(
f"Loaded model defaults from {config_path} (via mapping)"
)
return config
# If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from
# adapter_config.json), try matching the last 1-2 path components against
# the registry (e.g. "Spark-TTS-0.5B/LLM").
if model_name not in _REVERSE_MODEL_MAPPING and (model_name.startswith("/") or model_name.startswith(".")):
if model_name not in _REVERSE_MODEL_MAPPING and (
model_name.startswith("/") or model_name.startswith(".")
):
parts = Path(model_name).parts
for depth in [2, 1]:
if len(parts) >= depth:
@ -1263,9 +1320,11 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
canonical_file = _REVERSE_MODEL_MAPPING[suffix]
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path} (via path suffix '{suffix}')")
logger.info(
f"Loaded model defaults from {config_path} (via path suffix '{suffix}')"
)
return config
# Try exact model name match (for backward compatibility)
@ -1273,48 +1332,58 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
# Search in subfolders and root
for config_path in defaults_dir.rglob(model_filename):
if config_path.is_file():
with open(config_path, 'r', encoding='utf-8') as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path}")
return config
# Fall back to default.yaml
default_config_path = defaults_dir / "default.yaml"
if default_config_path.exists():
with open(default_config_path, 'r', encoding='utf-8') as f:
with open(default_config_path, "r", encoding = "utf-8") as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded default model defaults from {default_config_path}")
return config
logger.warning(f"No default config found for model {model_name}")
return {}
except Exception as e:
logger.error(f"Error loading model defaults for {model_name}: {e}")
return {}
@dataclass
class ModelConfig:
"""Configuration for a model to load"""
identifier: str # Clean model identifier (org/name or path)
display_name: str # Original UI display name
path: str # Normalized filesystem path
is_local: bool # Is this a local file vs HF model?
is_cached: bool # Is this already in HF cache?
is_vision: bool # Is this a vision model?
is_lora: bool # Is this a lora adapter?
is_gguf: bool = False # Is this a GGUF model?
is_audio: bool = False # Is this a TTS audio model?
audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
identifier: str # Clean model identifier (org/name or path)
display_name: str # Original UI display name
path: str # Normalized filesystem path
is_local: bool # Is this a local file vs HF model?
is_cached: bool # Is this already in HF cache?
is_vision: bool # Is this a vision model?
is_lora: bool # Is this a lora adapter?
is_gguf: bool = False # Is this a GGUF model?
is_audio: bool = False # Is this a TTS audio model?
audio_type: Optional[str] = (
None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
)
has_audio_input: bool = False # Accepts audio input (ASR/speech understanding)
gguf_file: Optional[str] = None # Full path to the .gguf file (local mode)
gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection)
gguf_hf_repo: Optional[str] = None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
gguf_mmproj_file: Optional[str] = (
None # Full path to the mmproj .gguf file (vision projection)
)
gguf_hf_repo: Optional[str] = (
None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
)
gguf_variant: Optional[str] = None # Quantization variant (e.g. "Q4_K_M")
base_model: Optional[str] = None # Base model (for LoRAs)
@classmethod
def from_lora_path(cls, lora_path: str, hf_token: Optional[str] = None) -> Optional['ModelConfig']:
def from_lora_path(
cls, lora_path: str, hf_token: Optional[str] = None
) -> Optional["ModelConfig"]:
"""
Create ModelConfig from a local LoRA adapter path.
@ -1341,26 +1410,26 @@ class ModelConfig:
return None
# Check if base model is vision
is_vision = is_vision_model(base_model, hf_token=hf_token)
is_vision = is_vision_model(base_model, hf_token = hf_token)
# Check if base model is audio
audio_type = detect_audio_type(base_model, hf_token=hf_token)
audio_type = detect_audio_type(base_model, hf_token = hf_token)
display_name = lora_path_obj.name
identifier = lora_path # Use path as identifier for local LoRAs
return cls(
identifier=identifier,
display_name=display_name,
path=lora_path,
is_local=True,
is_cached=True, # Local LoRAs are always "cached"
is_vision=is_vision,
is_lora=True,
is_audio=audio_type is not None and audio_type != 'audio_vlm',
audio_type=audio_type,
has_audio_input=is_audio_input_type(audio_type),
base_model=base_model,
identifier = identifier,
display_name = display_name,
path = lora_path,
is_local = True,
is_cached = True, # Local LoRAs are always "cached"
is_vision = is_vision,
is_lora = True,
is_audio = audio_type is not None and audio_type != "audio_vlm",
audio_type = audio_type,
has_audio_input = is_audio_input_type(audio_type),
base_model = base_model,
)
except Exception as e:
@ -1374,7 +1443,7 @@ class ModelConfig:
hf_token: Optional[str] = None,
is_lora: bool = False,
gguf_variant: Optional[str] = None,
) -> Optional['ModelConfig']:
) -> Optional["ModelConfig"]:
"""
Create ModelConfig from a clean model identifier.
@ -1432,7 +1501,7 @@ class ModelConfig:
try:
meta = json.loads(meta_path.read_text())
base = meta.get("base_model")
if base and is_vision_model(base, hf_token=hf_token):
if base and is_vision_model(base, hf_token = hf_token):
base_is_vision = True
logger.info(f"GGUF base model '{base}' is a vision model")
except Exception as e:
@ -1444,27 +1513,30 @@ class ModelConfig:
gguf_is_vision = True
logger.info(f"Detected mmproj for vision: {mmproj_file}")
elif base_is_vision:
logger.warning(f"Base model is vision but no mmproj file found in {gguf_dir}")
logger.warning(
f"Base model is vision but no mmproj file found in {gguf_dir}"
)
return cls(
identifier=identifier,
display_name=display_name,
path=path,
is_local=True,
is_cached=True,
is_vision=gguf_is_vision,
is_lora=False,
is_gguf=True,
gguf_file=gguf_file,
gguf_mmproj_file=mmproj_file,
identifier = identifier,
display_name = display_name,
path = path,
is_local = True,
is_cached = True,
is_vision = gguf_is_vision,
is_lora = False,
is_gguf = True,
gguf_file = gguf_file,
gguf_mmproj_file = mmproj_file,
)
else:
# Check if the HF repo contains GGUF files
gguf_filename = detect_gguf_model_remote(identifier, hf_token=hf_token)
gguf_filename = detect_gguf_model_remote(identifier, hf_token = hf_token)
if gguf_filename:
# Preflight: verify llama-server binary exists BEFORE user waits
# for a multi-GB download that llama-server handles natively
from core.inference.llama_cpp import LlamaCppBackend
if not LlamaCppBackend._find_llama_server_binary():
raise RuntimeError(
"llama-server binary not found — cannot load GGUF models. "
@ -1472,7 +1544,7 @@ class ModelConfig:
)
# Use list_gguf_variants() to detect vision & resolve variant
variants, has_vision = list_gguf_variants(identifier, hf_token=hf_token)
variants, has_vision = list_gguf_variants(identifier, hf_token = hf_token)
variant = gguf_variant
if not variant:
# Auto-select best quantization
@ -1489,17 +1561,17 @@ class ModelConfig:
f"variant={variant}, vision={has_vision}"
)
return cls(
identifier=identifier,
display_name=display_name,
path=identifier,
is_local=False,
is_cached=False,
is_vision=has_vision,
is_lora=False,
is_gguf=True,
gguf_file=None,
gguf_hf_repo=identifier,
gguf_variant=variant,
identifier = identifier,
display_name = display_name,
path = identifier,
is_local = False,
is_cached = False,
is_vision = has_vision,
is_lora = False,
is_gguf = True,
gguf_file = None,
gguf_hf_repo = identifier,
gguf_variant = variant,
)
# Auto-detect LoRA for local paths (check adapter_config.json on disk)
@ -1507,20 +1579,25 @@ class ModelConfig:
detected_base = get_base_model_from_lora(path)
if detected_base:
is_lora = True
logger.info(f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})")
logger.info(
f"Auto-detected local LoRA adapter at '{path}' (base: {detected_base})"
)
# Auto-detect LoRA for remote HF models (check repo file listing)
if not is_lora and not is_local:
try:
from huggingface_hub import model_info as hf_model_info
info = hf_model_info(identifier, token=hf_token)
info = hf_model_info(identifier, token = hf_token)
repo_files = [s.rfilename for s in info.siblings]
if "adapter_config.json" in repo_files:
is_lora = True
logger.info(f"Auto-detected remote LoRA adapter: '{identifier}'")
except Exception as e:
logger.debug(f"Could not check remote LoRA status for '{identifier}': {e}")
logger.debug(
f"Could not check remote LoRA status for '{identifier}': {e}"
)
# Handle LoRA adapters
base_model = None
if is_lora:
@ -1531,15 +1608,20 @@ class ModelConfig:
# Remote LoRA: download adapter_config.json from HF
try:
from huggingface_hub import hf_hub_download
config_path = hf_hub_download(identifier, "adapter_config.json", token=hf_token)
with open(config_path, 'r') as f:
config_path = hf_hub_download(
identifier, "adapter_config.json", token = hf_token
)
with open(config_path, "r") as f:
adapter_config = json.load(f)
base_model = adapter_config.get("base_model_name_or_path")
if base_model:
logger.info(f"Resolved remote LoRA base model: '{base_model}'")
except Exception as e:
logger.warning(f"Could not download adapter_config.json for '{identifier}': {e}")
logger.warning(
f"Could not download adapter_config.json for '{identifier}': {e}"
)
if not base_model:
logger.warning(f"Could not determine base model for LoRA '{path}'")
return None
@ -1547,34 +1629,35 @@ class ModelConfig:
else:
check_model = identifier
vision = is_vision_model(check_model, hf_token=hf_token)
audio_type_val = detect_audio_type(check_model, hf_token=hf_token)
vision = is_vision_model(check_model, hf_token = hf_token)
audio_type_val = detect_audio_type(check_model, hf_token = hf_token)
has_audio_in = is_audio_input_type(audio_type_val)
display_name = Path(path).name if is_local else identifier.split("/")[-1]
return cls(
identifier=identifier,
display_name=display_name,
path=path,
is_local=is_local,
is_cached=is_model_cached(identifier) if not is_local else True,
is_vision=vision,
is_lora=is_lora,
is_audio=audio_type_val is not None and audio_type_val != 'audio_vlm',
audio_type=audio_type_val,
has_audio_input=has_audio_in,
base_model=base_model,
identifier = identifier,
display_name = display_name,
path = path,
is_local = is_local,
is_cached = is_model_cached(identifier) if not is_local else True,
is_vision = vision,
is_lora = is_lora,
is_audio = audio_type_val is not None and audio_type_val != "audio_vlm",
audio_type = audio_type_val,
has_audio_input = has_audio_in,
base_model = base_model,
)
@classmethod
def from_ui_selection(cls,
dropdown_value: Optional[str],
search_value: Optional[str],
local_models: list = None,
hf_token: Optional[str] = None,
is_lora: bool = False) -> Optional['ModelConfig']:
def from_ui_selection(
cls,
dropdown_value: Optional[str],
search_value: Optional[str],
local_models: list = None,
hf_token: Optional[str] = None,
is_lora: bool = False,
) -> Optional["ModelConfig"]:
"""
Create a universal ModelConfig from UI dropdown/search selections.
Handles base models and LoRA adapters.
@ -1592,7 +1675,9 @@ class ModelConfig:
# Use the correct 'local_models' parameter to resolve display names
if " (Active)" in selected or " (Ready)" in selected:
clean_display_name = selected.replace(" (Active)", "").replace(" (Ready)", "")
clean_display_name = selected.replace(" (Active)", "").replace(
" (Ready)", ""
)
if local_models:
for local_display, local_path in local_models:
if local_display == clean_display_name:
@ -1621,25 +1706,28 @@ class ModelConfig:
# For a LoRA, we MUST find its base model.
base_model = get_base_model_from_lora(path)
if not base_model:
logger.warning(f"Could not determine base model for LoRA '{path}'. Cannot create config.")
return None # Cannot proceed without a base model
logger.warning(
f"Could not determine base model for LoRA '{path}'. Cannot create config."
)
return None # Cannot proceed without a base model
# A LoRA's vision capability is determined by its base model.
is_vision = is_vision_model(base_model, hf_token=hf_token)
is_vision = is_vision_model(base_model, hf_token = hf_token)
else:
# For a base model, just check its own vision status.
is_vision = is_vision_model(identifier, hf_token=hf_token)
is_vision = is_vision_model(identifier, hf_token = hf_token)
from utils.paths import is_model_cached
is_cached = is_model_cached(identifier) if not is_local else True
return cls(
identifier=identifier,
display_name=display_name,
path=path,
is_local=is_local,
is_cached=is_cached,
is_vision=is_vision,
is_lora=is_lora,
base_model=base_model, # This will be None for base models, and populated for LoRAs
identifier = identifier,
display_name = display_name,
path = path,
is_local = is_local,
is_cached = is_cached,
is_vision = is_vision,
is_lora = is_lora,
base_model = base_model, # This will be None for base models, and populated for LoRAs
)

51
studio/backend/utils/paths/__init__.py
View file

@ -4,6 +4,7 @@
"""
Path utilities for model and dataset handling
"""
from .path_utils import normalize_path, is_local_path, is_model_cached, get_cache_path
from .storage_roots import (
studio_root,
@ -21,6 +22,7 @@ from .storage_roots import (
oxc_validator_tmp_root,
tensorboard_root,
ensure_dir,
ensure_studio_directories,
resolve_under_root,
resolve_output_dir,
resolve_export_dir,
@ -29,28 +31,29 @@ from .storage_roots import (
)
__all__ = [
'normalize_path',
'is_local_path',
'is_model_cached',
'get_cache_path',
'studio_root',
'assets_root',
'datasets_root',
'dataset_uploads_root',
'recipe_datasets_root',
'outputs_root',
'exports_root',
'auth_root',
'auth_db_path',
'tmp_root',
'seed_uploads_root',
'unstructured_seed_cache_root',
'oxc_validator_tmp_root',
'tensorboard_root',
'ensure_dir',
'resolve_under_root',
'resolve_output_dir',
'resolve_export_dir',
'resolve_tensorboard_dir',
'resolve_dataset_path',
"normalize_path",
"is_local_path",
"is_model_cached",
"get_cache_path",
"studio_root",
"assets_root",
"datasets_root",
"dataset_uploads_root",
"recipe_datasets_root",
"outputs_root",
"exports_root",
"auth_root",
"auth_db_path",
"tmp_root",
"seed_uploads_root",
"unstructured_seed_cache_root",
"oxc_validator_tmp_root",
"tensorboard_root",
"ensure_dir",
"ensure_studio_directories",
"resolve_under_root",
"resolve_output_dir",
"resolve_export_dir",
"resolve_tensorboard_dir",
"resolve_dataset_path",
]

34
studio/backend/utils/paths/path_utils.py
View file

@ -4,6 +4,7 @@
"""
Path utilities for model and dataset handling
"""
import os
from pathlib import Path
from typing import Optional
@ -25,14 +26,14 @@ def normalize_path(path: str) -> str:
return path
# Handle Windows drive letters (C:\\ or c:\\)
if len(path) >= 3 and path[1] == ':' and path[2] in ('\\', '/'):
if len(path) >= 3 and path[1] == ":" and path[2] in ("\\", "/"):
drive = path[0].lower()
rest = path[3:].replace('\\', '/')
return f'/mnt/{drive}/{rest}'
rest = path[3:].replace("\\", "/")
return f"/mnt/{drive}/{rest}"
# Already Unix-style or relative
return path.replace('\\', '/')
pass
return path.replace("\\", "/")
def is_local_path(path: str) -> bool:
"""
@ -53,26 +54,26 @@ def is_local_path(path: str) -> bool:
pass
# Obvious HF patterns
if path.count('/') == 1 and not path.startswith(('/', '.', '~')):
if path.count("/") == 1 and not path.startswith(("/", ".", "~")):
return False # Looks like org/model format
# Filesystem indicators
return (
path.startswith(('/', '.', '~')) or # Unix absolute/relative
':' in path or # Windows drive or URL
'\\' in path or # Windows separator
os.path.isabs(path) # System-absolute
path.startswith(("/", ".", "~")) # Unix absolute/relative
or ":" in path # Windows drive or URL
or "\\" in path # Windows separator
or os.path.isabs(path) # System-absolute
)
pass
def get_cache_path(model_name: str) -> Optional[Path]:
"""Get HuggingFace cache path for a model if it exists."""
cache_dir = Path.home() / '.cache' / 'huggingface' / 'hub'
cache_dir = Path.home() / ".cache" / "huggingface" / "hub"
model_cache_name = model_name.replace("/", "--")
model_cache_path = cache_dir / f'models--{model_cache_name}'
model_cache_path = cache_dir / f"models--{model_cache_name}"
return model_cache_path if model_cache_path.exists() else None
pass
def is_model_cached(model_name: str) -> bool:
"""Check if model is downloaded in HuggingFace cache."""
@ -81,9 +82,8 @@ def is_model_cached(model_name: str) -> bool:
return False
# Check for actual model files
for suffix in ['.safetensors', '.bin', '.json']:
if list(cache_path.rglob(f'*{suffix}')):
for suffix in [".safetensors", ".bin", ".json"]:
if list(cache_path.rglob(f"*{suffix}")):
return True
return False
pass

39
studio/backend/utils/paths/storage_roots.py
View file

@ -1,3 +1,6 @@
# 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
from pathlib import Path
@ -61,11 +64,29 @@ def tensorboard_root() -> Path:
def ensure_dir(path: Path) -> Path:
path.mkdir(parents=True, exist_ok=True)
path.mkdir(parents = True, exist_ok = True)
return path
def _clean_relative_path(path_value: str, *, strip_prefixes: tuple[str, ...] = ()) -> Path:
def ensure_studio_directories() -> None:
"""Create all standard studio directories on startup."""
for dir_fn in (
studio_root,
assets_root,
datasets_root,
dataset_uploads_root,
recipe_datasets_root,
outputs_root,
exports_root,
auth_root,
tensorboard_root,
):
ensure_dir(dir_fn())
def _clean_relative_path(
path_value: str, *, strip_prefixes: tuple[str, ...] = ()
) -> Path:
path = Path(path_value).expanduser()
parts = [part for part in path.parts if part not in ("", ".")]
while parts and parts[0] in strip_prefixes:
@ -86,31 +107,31 @@ def resolve_under_root(
if path.is_absolute():
return path
cleaned = _clean_relative_path(str(path), strip_prefixes=strip_prefixes)
cleaned = _clean_relative_path(str(path), strip_prefixes = strip_prefixes)
return root / cleaned
def resolve_output_dir(path_value: str | None = None) -> Path:
return resolve_under_root(
path_value,
root=outputs_root(),
strip_prefixes=("outputs",),
root = outputs_root(),
strip_prefixes = ("outputs",),
)
def resolve_export_dir(path_value: str | None = None) -> Path:
return resolve_under_root(
path_value,
root=exports_root(),
strip_prefixes=("exports",),
root = exports_root(),
strip_prefixes = ("exports",),
)
def resolve_tensorboard_dir(path_value: str | None = None) -> Path:
return resolve_under_root(
path_value,
root=tensorboard_root(),
strip_prefixes=("runs", "tensorboard"),
root = tensorboard_root(),
strip_prefixes = ("runs", "tensorboard"),
)

59
studio/backend/utils/transformers_version.py
View file

@ -33,7 +33,6 @@ from pathlib import Path
logger = get_logger(__name__)
# ---------------------------------------------------------------------------
# Detection
# ---------------------------------------------------------------------------
@ -41,12 +40,12 @@ logger = get_logger(__name__)
# Lowercase substrings — if ANY appears anywhere in the lowered model name,
# we need transformers 5.x.
TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
"ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512
"glm-4.7-flash", # GLM-4.7-Flash
"qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants
"qwen3.5", # Qwen3.5 family (35B-A3B, etc.)
"qwen3-next", # Qwen3-Next and variants
"tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B
"ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512
"glm-4.7-flash", # GLM-4.7-Flash
"qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants
"qwen3.5", # Qwen3.5 family (35B-A3B, etc.)
"qwen3-next", # Qwen3-Next and variants
"tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B
)
# Versions
@ -77,7 +76,8 @@ def _resolve_base_model(model_name: str) -> str:
if base:
logger.info(
"Resolved LoRA adapter '%s' → base model '%s'",
model_name, base,
model_name,
base,
)
return base
except Exception as exc:
@ -87,18 +87,21 @@ def _resolve_base_model(model_name: str) -> str:
if local_path.is_dir():
try:
from utils.models import get_base_model_from_lora
base = get_base_model_from_lora(model_name)
if base:
logger.info(
"Resolved LoRA adapter '%s' → base model '%s' "
"(via get_base_model_from_lora)",
model_name, base,
model_name,
base,
)
return base
except Exception as exc:
logger.debug(
"get_base_model_from_lora failed for '%s': %s",
model_name, exc,
model_name,
exc,
)
return model_name
@ -115,6 +118,7 @@ def needs_transformers_5(model_name: str) -> bool:
# Version switching (in-process — used only by export)
# ---------------------------------------------------------------------------
def _get_in_memory_version() -> str | None:
"""Return the transformers version currently loaded in this process."""
tf = sys.modules.get("transformers")
@ -153,7 +157,8 @@ def _purge_modules() -> int:
"""
importlib.invalidate_caches()
to_remove = [
k for k in list(sys.modules.keys())
k
for k in list(sys.modules.keys())
if any(k == p or k.startswith(p + ".") for p in _PURGE_PREFIXES)
]
for key in to_remove:
@ -167,15 +172,21 @@ def _ensure_venv_t5_exists() -> bool:
return True
logger.warning(".venv_t5 not found at %s — installing at runtime", _VENV_T5_DIR)
os.makedirs(_VENV_T5_DIR, exist_ok=True)
os.makedirs(_VENV_T5_DIR, exist_ok = True)
for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub==1.3.0"):
cmd = [
sys.executable, "-m", "pip", "install",
"--target", _VENV_T5_DIR,
sys.executable,
"-m",
"pip",
"install",
"--target",
_VENV_T5_DIR,
"--no-deps",
pkg,
]
result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
result = subprocess.run(
cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, text = True
)
if result.returncode != 0:
logger.error("pip install failed:\n%s", result.stdout)
return False
@ -186,7 +197,9 @@ def _ensure_venv_t5_exists() -> bool:
def _activate_5x() -> None:
"""Prepend .venv_t5/ to sys.path, purge stale modules, reimport."""
if not _ensure_venv_t5_exists():
raise RuntimeError(f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}")
raise RuntimeError(
f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}"
)
if _VENV_T5_DIR not in sys.path:
sys.path.insert(0, _VENV_T5_DIR)
@ -196,6 +209,7 @@ def _activate_5x() -> None:
logger.info("Purged %d cached modules", count)
import transformers
logger.info("Loaded transformers %s", transformers.__version__)
@ -209,6 +223,7 @@ def _deactivate_5x() -> None:
logger.info("Purged %d cached modules", count)
import transformers
logger.info("Reverted to transformers %s", transformers.__version__)
@ -236,7 +251,10 @@ def ensure_transformers_version(model_name: str) -> None:
logger.info(
"Version check for '%s' (resolved: '%s'): need=%s, in_memory=%s",
model_name, resolved, target_version, in_memory,
model_name,
resolved,
target_version,
in_memory,
)
# --- Already correct? ---------------------------------------------------
@ -245,7 +263,8 @@ def ensure_transformers_version(model_name: str) -> None:
if in_memory_major == target_major:
logger.info(
"transformers %s already loaded — correct for '%s'",
in_memory, model_name,
in_memory,
model_name,
)
return
@ -254,7 +273,9 @@ def ensure_transformers_version(model_name: str) -> None:
logger.info("Activating transformers %s via .venv_t5…", TRANSFORMERS_5_VERSION)
_activate_5x()
else:
logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION)
logger.info(
"Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION
)
_deactivate_5x()
final = _get_in_memory_version()

36
studio/backend/utils/utils.py
View file

@ -4,6 +4,7 @@
"""
Shared backend utilities
"""
import os
import structlog
from loggers import get_logger
@ -28,26 +29,26 @@ def without_hf_auth():
"""
# Save environment variables
saved_env = {}
env_vars = ['HF_TOKEN', 'HUGGINGFACE_HUB_TOKEN', 'HF_HOME']
env_vars = ["HF_TOKEN", "HUGGINGFACE_HUB_TOKEN", "HF_HOME"]
for var in env_vars:
if var in os.environ:
saved_env[var] = os.environ[var]
del os.environ[var]
# Save disable flag
saved_disable = os.environ.get('HF_HUB_DISABLE_IMPLICIT_TOKEN')
os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = '1'
saved_disable = os.environ.get("HF_HUB_DISABLE_IMPLICIT_TOKEN")
os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "1"
# Move token files temporarily
token_files = []
token_locations = [
Path.home() / '.cache' / 'huggingface' / 'token',
Path.home() / '.huggingface' / 'token'
Path.home() / ".cache" / "huggingface" / "token",
Path.home() / ".huggingface" / "token",
]
for token_loc in token_locations:
if token_loc.exists():
temp = tempfile.NamedTemporaryFile(delete=False)
temp = tempfile.NamedTemporaryFile(delete = False)
temp.close()
shutil.move(str(token_loc), temp.name)
token_files.append((token_loc, temp.name))
@ -58,7 +59,7 @@ def without_hf_auth():
# Restore tokens
for original, temp in token_files:
try:
original.parent.mkdir(parents=True, exist_ok=True)
original.parent.mkdir(parents = True, exist_ok = True)
shutil.move(temp, str(original))
except Exception as e:
logger.error(f"Failed to restore token {original}: {e}")
@ -68,10 +69,10 @@ def without_hf_auth():
os.environ[var] = value
if saved_disable is not None:
os.environ['HF_HUB_DISABLE_IMPLICIT_TOKEN'] = saved_disable
os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = saved_disable
else:
os.environ.pop('HF_HUB_DISABLE_IMPLICIT_TOKEN', None)
pass
os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None)
def format_error_message(error: Exception, model_name: str) -> str:
"""
@ -85,7 +86,7 @@ def format_error_message(error: Exception, model_name: str) -> str:
User-friendly error string
"""
error_str = str(error).lower()
model_short = model_name.split('/')[-1] if '/' in model_name else model_name
model_short = model_name.split("/")[-1] if "/" in model_name else model_name
if "repository not found" in error_str or "404" in error_str:
return f"Model '{model_short}' not found. Check the model name."
@ -99,12 +100,19 @@ def format_error_message(error: Exception, model_name: str) -> str:
if "invalid user token" in error_str:
return "Invalid HF token. Please check your token and try again."
if "memory" in error_str or "cuda" in error_str or "mlx" in error_str or "out of memory" in error_str:
if (
"memory" in error_str
or "cuda" in error_str
or "mlx" in error_str
or "out of memory" in error_str
):
from utils.hardware import get_device
device = get_device()
device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(device.value, "GPU")
device_label = {"cuda": "GPU", "mlx": "Apple Silicon GPU", "cpu": "system"}.get(
device.value, "GPU"
)
return f"Not enough {device_label} memory to load '{model_short}'. Try a smaller model or free memory."
# Generic fallback
return str(error)
pass