unslothai--unsloth/studio/backend/hub/utils/dataset_format.py

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

from __future__ import annotations

import re
from typing import Any, Optional


def _first_row(dataset) -> Optional[dict]:
    try:
        row = next(iter(dataset))
    except StopIteration:
        return None
    return row if isinstance(row, dict) else None


def _column_names(dataset, sample: Optional[dict] = None) -> list[str]:
    names = getattr(dataset, "column_names", None)
    if names is not None:
        return list(names)
    return list((sample or {}).keys())


def _keyword_in_column(keyword: str, col_name: str) -> bool:
    return re.search(r"\b" + re.escape(keyword) + r"\b", col_name, re.IGNORECASE) is not None


def _unknown_dataset_format(
    chat_column: Optional[str] = None, sample_keys: Optional[list[str]] = None
) -> dict:
    return {
        "format": "unknown",
        "chat_column": chat_column,
        "needs_standardization": None,
        "sample_keys": sample_keys or [],
    }


def detect_dataset_format(dataset) -> dict:
    sample = _first_row(dataset)
    if sample is None:
        return _unknown_dataset_format()
    column_names = set(sample.keys())
    if {"instruction", "output"}.issubset(column_names):
        return {
            "format": "alpaca",
            "chat_column": None,
            "needs_standardization": False,
            "sample_keys": [],
        }

    chat_column = None
    if "messages" in column_names:
        chat_column = "messages"
    elif "conversations" in column_names:
        chat_column = "conversations"
    elif "texts" in column_names:
        chat_column = "texts"

    if not chat_column:
        return _unknown_dataset_format()

    chat_data = sample.get(chat_column)
    if not isinstance(chat_data, (list, tuple)) or not chat_data:
        return _unknown_dataset_format(chat_column)
    first_msg = chat_data[0]
    if not isinstance(first_msg, dict):
        return _unknown_dataset_format(chat_column)
    msg_keys = set(first_msg.keys())
    sample_keys = [str(key) for key in msg_keys]
    if "from" in msg_keys or "value" in msg_keys:
        return {
            "format": "sharegpt",
            "chat_column": chat_column,
            "needs_standardization": True,
            "sample_keys": sample_keys,
        }
    if "role" in msg_keys and "content" in msg_keys:
        return {
            "format": "chatml",
            "chat_column": chat_column,
            "needs_standardization": False,
            "sample_keys": sample_keys,
        }
    return _unknown_dataset_format(chat_column, sample_keys)


def detect_custom_format_heuristic(dataset):
    sample = _first_row(dataset)
    if sample is None:
        return None
    all_columns = list(sample.keys())
    mapping = {}
    assistant_words = [
        "output",
        "answer",
        "response",
        "assistant",
        "completion",
        "expected",
        "recommendation",
        "reply",
        "result",
        "target",
        "solution",
        "explanation",
        "solve",
    ]
    user_words_high_priority = [
        "input",
        "question",
        "query",
        "prompt",
        "instruction",
        "request",
        "snippet",
        "user",
        "text",
        "problem",
        "exercise",
    ]
    user_words_low_priority = ["task"]
    user_words = user_words_high_priority + user_words_low_priority
    system_words = [
        "system",
        "context",
        "description",
        "persona",
        "role",
        "template",
        "task",
    ]
    metadata_exact_match = {
        "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",
    ]
    priority_patterns = {
        "generated": 100,
        "gen_": 90,
        "model_": 80,
        "predicted": 70,
        "completion": 60,
    }

    def has_keyword(col_name, keywords):
        col_lower = col_name.lower()
        col_normalized = col_lower.replace("_", "").replace("-", "").replace(" ", "")
        return any(keyword in col_lower or keyword in col_normalized for keyword in keywords)

    def is_metadata(col_name):
        col_lower = col_name.lower()
        if col_lower in metadata_exact_match or col_lower in metadata_prefix_patterns:
            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"]:
                        return True
        return len(col_lower) <= 2 and col_lower not in ["qa", "q", "a"]

    def get_priority_score(col_name):
        col_lower = col_name.lower()
        return sum(score for pattern, score in priority_patterns.items() if pattern in col_lower)

    def get_content_length(col_name):
        try:
            return len(str(sample[col_name])) if sample.get(col_name) else 0
        except Exception:
            return 0

    def score_column(col_name, keywords, role_type, num_candidates):
        if not has_keyword(col_name, keywords):
            return 0
        score = 10
        if role_type == "user":
            col_lower = col_name.lower()
            if "task" in col_lower and not any(kw in col_lower for kw in user_words_high_priority):
                score -= 15
        score += get_priority_score(col_name)
        if role_type in ["assistant", "user"]:
            avg_length = get_content_length(col_name)
            if num_candidates > 1:
                if avg_length > 1000:
                    score += 50
                elif avg_length > 200:
                    score += 30
                elif avg_length > 50:
                    score += 10
                elif avg_length < 50:
                    score -= 20
            else:
                if avg_length > 1000:
                    score += 50
                elif avg_length > 200:
                    score += 30
                elif avg_length > 50:
                    score += 10
        return score

    content_columns = [col for col in all_columns if not is_metadata(col)]
    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)]
    assistant_candidates = [
        (col, score)
        for col in assistant_potential
        if (score := score_column(col, assistant_words, "assistant", len(assistant_potential))) > 0
    ]
    if assistant_candidates:
        assistant_candidates.sort(key = lambda item: item[1], reverse = True)
        assistant_col = assistant_candidates[0][0]
        mapping[assistant_col] = "assistant"
    else:
        assistant_col = None

    user_candidates = []
    for col in user_potential:
        if col == assistant_col:
            continue
        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 item: item[1], reverse = True)
        user_col = user_candidates[0][0]
        mapping[user_col] = "user"
    else:
        user_col = None

    remaining_columns = [col for col in content_columns if col not in mapping]
    system_col = None
    for col in remaining_columns:
        if has_keyword(col, system_words):
            mapping[col] = "system"
            system_col = col
            break
    if system_col:
        remaining_columns = [col for col in remaining_columns if col != system_col]
    if remaining_columns:
        remaining_col = remaining_columns[0]
        if not has_keyword(remaining_col, user_words + assistant_words):
            mapping[remaining_col] = "system"
        elif user_col is None:
            mapping[remaining_col] = "user"
        else:
            mapping[remaining_col] = "system"

    has_user = any(role == "user" for role in mapping.values())
    has_assistant = any(role == "assistant" for role in mapping.values())
    if not has_user:
        for col in remaining_columns:
            if col not in mapping:
                mapping[col] = "user"
                has_user = True
                break
    return mapping if has_user and has_assistant else None


_AUDIO_EXTENSIONS = (
    ".wav",
    ".mp3",
    ".flac",
    ".ogg",
    ".opus",
    ".m4a",
    ".aac",
    ".wma",
    ".webm",
)


def _is_audio_value(value) -> bool:
    if value is None:
        return False
    if isinstance(value, dict):
        if "array" in value and "sampling_rate" in value:
            return True
        if "bytes" in value or "path" in value:
            path = value.get("path") or ""
            return isinstance(path, str) and any(
                path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
            )
    return False


def _has_image_header(data: bytes) -> bool:
    if len(data) < 4:
        return False
    return (
        data[:2] == b"\xff\xd8"
        or data[:4] == b"\x89PNG"
        or data[:3] == b"GIF"
        or (data[:4] == b"RIFF" and len(data) >= 12 and data[8:12] == b"WEBP")
        or data[:2] == b"BM"
    )


def _is_image_value(value) -> bool:
    if value is None:
        return False
    try:
        from PIL.Image import Image as PILImage
        if isinstance(value, PILImage):
            return True
    except ImportError:
        pass
    if isinstance(value, dict):
        if "array" in value and "sampling_rate" in value:
            return False
        if "bytes" in value and "path" in value:
            path = value.get("path") or ""
            if isinstance(path, str) and any(
                path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS
            ):
                return False
            return True
    if isinstance(value, (bytes, bytearray)):
        return _has_image_header(value)
    image_exts = (".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tiff", ".svg")
    if isinstance(value, str) and len(value) < 1000:
        lower = value.strip().lower()
        if lower.startswith(("http://", "https://")):
            return any(lower.split("?")[0].endswith(ext) for ext in image_exts)
        return any(lower.endswith(ext) for ext in image_exts)
    return False


def detect_multimodal_dataset(dataset):
    sample = _first_row(dataset)
    if sample is None:
        return {
            "is_image": False,
            "multimodal_columns": [],
            "modality_types": [],
            "is_audio": False,
            "audio_columns": [],
            "detected_audio_column": None,
            "detected_text_column": None,
            "detected_speaker_column": None,
        }
    column_names = list(sample.keys())
    image_keywords = [
        "image",
        "img",
        "pixel",
        "jpg",
        "jpeg",
        "png",
        "webp",
        "bmp",
        "gif",
        "tiff",
        "svg",
        "photo",
        "pic",
        "picture",
        "visual",
        "file_name",
        "filename",
    ]
    audio_keywords = ["audio", "speech", "wav", "waveform", "sound"]
    multimodal_columns = []
    audio_columns = []
    modality_types = set()
    for col_name in column_names:
        if any(_keyword_in_column(keyword, col_name) for keyword in image_keywords):
            multimodal_columns.append(col_name)
            modality_types.add("image")
    for col_name in column_names:
        if col_name not in multimodal_columns and _is_image_value(sample[col_name]):
            multimodal_columns.append(col_name)
            modality_types.add("image")
    for col_name in column_names:
        if any(_keyword_in_column(keyword, col_name) for keyword in audio_keywords):
            audio_columns.append(col_name)
            modality_types.add("audio")
    for col_name in column_names:
        if col_name not in audio_columns and _is_audio_value(sample[col_name]):
            audio_columns.append(col_name)
            modality_types.add("audio")
    if audio_columns:
        multimodal_columns = [col for col in multimodal_columns if col not in set(audio_columns)]

    detected_text_col = None
    if audio_columns:
        for col_name in column_names:
            if col_name.lower() in [
                "text",
                "sentence",
                "transcript",
                "transcription",
                "label",
            ]:
                detected_text_col = col_name
                break
    detected_speaker_col = None
    if audio_columns:
        for col_name in column_names:
            if col_name.lower() in ["source", "speaker", "speaker_id"]:
                detected_speaker_col = col_name
                break
    return {
        "is_image": len(multimodal_columns) > 0,
        "multimodal_columns": multimodal_columns,
        "modality_types": list(modality_types),
        "is_audio": len(audio_columns) > 0,
        "audio_columns": audio_columns,
        "detected_audio_column": audio_columns[0] if audio_columns else None,
        "detected_text_column": detected_text_col,
        "detected_speaker_column": detected_speaker_col,
    }


def detect_vlm_dataset_structure(dataset):
    sample = _first_row(dataset)
    if sample is None:
        return {
            "format": "unknown",
            "needs_conversion": None,
            "image_column": None,
            "text_column": None,
            "messages_column": None,
        }
    column_names = set(sample.keys())
    if "messages" in column_names:
        messages = sample["messages"]
        if messages and len(messages) > 0:
            first_msg = messages[0]
            if "content" in first_msg:
                content = first_msg["content"]
                if (
                    isinstance(content, list)
                    and content
                    and isinstance(content[0], dict)
                    and "type" in content[0]
                ):
                    has_index = any("index" in item for item in content if isinstance(item, dict))
                    if has_index and "images" in column_names:
                        return {
                            "format": "vlm_messages_llava",
                            "needs_conversion": True,
                            "messages_column": "messages",
                            "image_column": "images",
                            "text_column": None,
                        }
                    has_image = any("image" in item for item in content if isinstance(item, dict))
                    if has_image:
                        return {
                            "format": "vlm_messages",
                            "needs_conversion": False,
                            "messages_column": "messages",
                            "image_column": None,
                            "text_column": None,
                        }

    for chat_col in ("conversations", "messages"):
        if chat_col not in column_names:
            continue
        chat_data = sample[chat_col]
        if not isinstance(chat_data, list) or not chat_data:
            continue
        has_image_placeholder = any(
            "<image>" in str(message.get("value", "") or message.get("content", ""))
            for message in chat_data
            if isinstance(message, dict)
        )
        if not has_image_placeholder:
            continue
        image_col = next(
            (
                col
                for col in column_names
                if col != chat_col
                and (_keyword_in_column("image", col) or _keyword_in_column("img", col))
            ),
            None,
        )
        if image_col:
            return {
                "format": "sharegpt_with_images",
                "needs_conversion": True,
                "image_column": image_col,
                "text_column": None,
                "messages_column": chat_col,
            }

    metadata_suffixes = (
        "_id",
        "_url",
        "_name",
        "_filename",
        "_uri",
        "_link",
        "_key",
        "_index",
    )
    metadata_prefixes = (
        "id_",
        "url_",
        "name_",
        "filename_",
        "uri_",
        "link_",
        "key_",
        "index_",
    )
    image_keywords = [
        "image",
        "img",
        "photo",
        "picture",
        "pic",
        "visual",
        "scan",
        "file_name",
        "filename",
    ]
    text_keywords = [
        "text",
        "caption",
        "captions",
        "description",
        "answer",
        "output",
        "response",
        "label",
    ]

    def is_metadata_column(col_name):
        lower = col_name.lower()
        return any(lower.endswith(suffix) for suffix in metadata_suffixes) or any(
            lower.startswith(prefix) for prefix in metadata_prefixes
        )

    image_candidates = []
    for col in column_names:
        value = sample[col]
        if any(_keyword_in_column(keyword, col) for keyword in image_keywords) or _is_image_value(
            value
        ):
            if hasattr(value, "size") and hasattr(value, "mode"):
                score = 100
            elif isinstance(value, dict) and ("bytes" in value or "path" in value):
                score = 75
            elif isinstance(value, str):
                score = (
                    55
                    if is_metadata_column(col)
                    else 70
                    if value.startswith(("http://", "https://"))
                    else 50
                )
            else:
                score = 0
            if score > 0:
                image_candidates.append((col, score))
    image_candidates.sort(key = lambda item: item[1], reverse = True)

    text_candidates = []
    for col in column_names:
        if is_metadata_column(col) or not any(
            _keyword_in_column(keyword, col) for keyword in text_keywords
        ):
            continue
        value = sample[col]
        if isinstance(value, str) and value:
            text_candidates.append((col, min(len(value), 1000)))
        elif isinstance(value, list) and value and isinstance(value[0], str):
            text_candidates.append((col, min(len(value[0]), 1000) // 2))
    text_candidates.sort(key = lambda item: item[1], reverse = True)

    found_image = image_candidates[0][0] if image_candidates else None
    found_text = text_candidates[0][0] if text_candidates else None
    if found_image and found_text:
        return {
            "format": "simple_image_text",
            "needs_conversion": True,
            "image_column": found_image,
            "text_column": found_text,
            "messages_column": None,
        }
    return {
        "format": "unknown",
        "needs_conversion": None,
        "image_column": found_image,
        "text_column": found_text,
        "messages_column": None,
    }


def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
    sample = _first_row(dataset)
    columns = _column_names(dataset, sample)
    multimodal_info = detect_multimodal_dataset(dataset)
    is_audio = multimodal_info.get("is_audio", False)
    audio_fields = {
        "is_audio": is_audio,
        "detected_audio_column": multimodal_info.get("detected_audio_column"),
        "detected_speaker_column": multimodal_info.get("detected_speaker_column"),
    }

    if is_vlm:
        vlm_structure = detect_vlm_dataset_structure(dataset)
        requires_mapping = vlm_structure["format"] == "unknown"
        warning = None
        if requires_mapping:
            missing = []
            if not vlm_structure.get("image_column"):
                missing.append("image")
            if not vlm_structure.get("text_column"):
                missing.append("text")
            if missing:
                warning = (
                    f"Could not auto-detect {' or '.join(missing)} column. "
                    "Please assign image and text columns manually."
                )
        return {
            "requires_manual_mapping": requires_mapping,
            "detected_format": vlm_structure["format"],
            "columns": columns,
            "suggested_mapping": None,
            "detected_image_column": vlm_structure.get("image_column"),
            "detected_text_column": vlm_structure.get("text_column"),
            "is_image": multimodal_info["is_image"],
            "multimodal_columns": multimodal_info.get("multimodal_columns"),
            "warning": warning,
            **audio_fields,
        }

    if is_audio:
        detected_audio = multimodal_info.get("detected_audio_column")
        detected_text = multimodal_info.get("detected_text_column")
        return {
            "requires_manual_mapping": not detected_audio or not detected_text,
            "detected_format": "audio",
            "columns": columns,
            "suggested_mapping": None,
            "detected_image_column": None,
            "detected_text_column": detected_text,
            "is_image": False,
            "multimodal_columns": multimodal_info.get("audio_columns"),
            **audio_fields,
        }

    detected = detect_dataset_format(dataset)
    if detected["format"] == "unknown":
        heuristic_mapping = detect_custom_format_heuristic(dataset)
        if heuristic_mapping:
            return {
                "requires_manual_mapping": False,
                "detected_format": "custom_heuristic",
                "columns": columns,
                "suggested_mapping": heuristic_mapping,
                "detected_image_column": None,
                "detected_text_column": None,
                "is_image": multimodal_info["is_image"],
                "multimodal_columns": multimodal_info.get("multimodal_columns"),
                **audio_fields,
            }
        return {
            "requires_manual_mapping": True,
            "detected_format": "unknown",
            "columns": columns,
            "suggested_mapping": None,
            "detected_image_column": None,
            "detected_text_column": None,
            "is_image": multimodal_info["is_image"],
            "multimodal_columns": multimodal_info.get("multimodal_columns"),
            "warning": (
                f"Could not auto-detect column roles for columns: {columns}. "
                "Please assign roles manually, or use AI Assist."
            ),
            **audio_fields,
        }

    return {
        "requires_manual_mapping": False,
        "detected_format": detected["format"],
        "columns": columns,
        "suggested_mapping": None,
        "detected_image_column": None,
        "detected_text_column": None,
        "is_image": multimodal_info["is_image"],
        "multimodal_columns": multimodal_info.get("multimodal_columns"),
        **audio_fields,
    }


_ROLE_MAP = {
    "human": "user",
    "user": "user",
    "input": "user",
    "gpt": "assistant",
    "assistant": "assistant",
    "output": "assistant",
    "system": "system",
}


def _standardize_sharegpt_row(row: dict[str, Any], chat_column: str) -> dict[str, Any]:
    chat_data = row.get(chat_column)
    if not isinstance(chat_data, list):
        return row
    messages = []
    for message in chat_data:
        if not isinstance(message, dict):
            continue
        role = message.get("role") or message.get("from")
        content = message.get("content") if "content" in message else message.get("value")
        messages.append(
            {
                "role": _ROLE_MAP.get(str(role), str(role or "user")),
                "content": "" if content is None else content,
            }
        )
    return {chat_column: messages}


def format_dataset_preview(dataset):
    detected = detect_dataset_format(dataset)
    if detected.get("format") != "sharegpt":
        return dataset
    chat_column = detected.get("chat_column")
    if not isinstance(chat_column, str):
        return dataset

    if hasattr(dataset, "map"):
        return dataset.map(lambda row: _standardize_sharegpt_row(row, chat_column))
    return dataset