YG_FT/backend/app/modules/data_process/dataset_format.py

"""Dataset format validation for Alpaca, ShareGPT, DPO, CPT formats.

Used by the training preflight flow to validate that uploaded dataset files
conform to the declared format before submitting to the compute node.
"""

from __future__ import annotations

import json
from typing import Any


def _load_sample(path: str | None, content: str | None = None, max_samples: int = 20) -> list[dict[str, Any]]:
    """Load up to max_samples records from JSONL file path or raw content string."""
    try:
        if content is not None:
            text = content.strip()
        elif path:
            with open(path, "r", encoding="utf-8") as fh:
                text = fh.read().strip()
        else:
            return []
    except Exception:
        return []

    if not text:
        return []

    # 先按整文件 JSON（数组/单对象）解析，兼容 .json；失败再按 jsonl 逐行解析
    try:
        value = json.loads(text)
    except json.JSONDecodeError:
        value = None
    if isinstance(value, list):
        return [item for item in value[:max_samples] if isinstance(item, dict)]
    if isinstance(value, dict):
        return [value]

    lines = text.splitlines()[:max_samples]
    records: list[dict[str, Any]] = []
    for line in lines:
        line = line.strip()
        if not line:
            continue
        try:
            record = json.loads(line)
        except json.JSONDecodeError:
            continue
        if isinstance(record, dict):
            records.append(record)
    return records


def _check_alpaca(records: list[dict[str, Any]]) -> list[str]:
    """Validate Alpaca format: requires 'instruction' field."""
    errors: list[str] = []
    if not records:
        errors.append("Alpaca 格式数据集无有效记录")
        return errors
    missing_instruction = sum(1 for r in records if not r.get("instruction"))
    if missing_instruction:
        errors.append(
            f"Alpaca 格式要求每条记录包含 instruction 字段，"
            f"前{len(records)}条中有{missing_instruction}条缺失"
        )
    return errors


def _check_sharegpt(records: list[dict[str, Any]]) -> list[str]:
    """Validate ShareGPT format: requires 'messages' (list of dicts with role/content)."""
    errors: list[str] = []
    if not records:
        errors.append("ShareGPT 格式数据集无有效记录")
        return errors
    bad = 0
    for r in records:
        messages = r.get("messages")
        if not isinstance(messages, list) or not messages:
            bad += 1
            continue
        for msg in messages:
            if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
                bad += 1
                break
    if bad:
        errors.append(
            f"ShareGPT 格式要求每条记录包含 messages 列表，"
            f"每条消息需有 role 和 content 字段，前{len(records)}条中有{bad}条不符合"
        )
    return errors


def _check_dpo(records: list[dict[str, Any]]) -> list[str]:
    """Validate DPO format: requires 'chosen' and 'rejected' fields."""
    errors: list[str] = []
    if not records:
        errors.append("DPO 格式数据集无有效记录")
        return errors
    missing_chosen = sum(1 for r in records if not r.get("chosen"))
    missing_rejected = sum(1 for r in records if not r.get("rejected"))
    if missing_chosen:
        errors.append(f"DPO 格式要求 chosen 字段，前{len(records)}条中有{missing_chosen}条缺失")
    if missing_rejected:
        errors.append(f"DPO 格式要求 rejected 字段，前{len(records)}条中有{missing_rejected}条缺失")
    return errors


def _check_cpt(records: list[dict[str, Any]]) -> list[str]:
    """Validate CPT format: requires 'text' field, should NOT have instruction/output."""
    errors: list[str] = []
    if not records:
        errors.append("CPT 格式数据集无有效记录")
        return errors
    missing_text = sum(1 for r in records if not r.get("text"))
    has_instruction = sum(1 for r in records if r.get("instruction") or r.get("output"))
    if missing_text:
        errors.append(f"CPT 格式要求 text 字段，前{len(records)}条中有{missing_text}条缺失")
    if has_instruction:
        errors.append(
            f"CPT 格式不应包含 instruction/output 字段（疑似 Alpaca 格式），"
            f"前{len(records)}条中有{has_instruction}条包含此类字段"
        )
    return errors


FORMAT_VALIDATORS = {
    "alpaca": _check_alpaca,
    "alpaca_jsonl": _check_alpaca,
    "sharegpt": _check_sharegpt,
    "dpo": _check_dpo,
    "cpt": _check_cpt,
    "pt": _check_cpt,
}


def validate_dataset_format(
    dataset_format: str,
    content: str | None = None,
    path: str | None = None,
    max_samples: int = 20,
) -> list[str]:
    """Validate dataset content against expected format.

    Args:
        dataset_format: One of 'alpaca', 'sharegpt', 'dpo', 'cpt'.
        content: Raw file content (JSONL text). Mutually exclusive with path.
        path: File path to read content from.
        max_samples: Maximum records to sample for validation.

    Returns:
        List of error messages (empty if valid).
    """
    fmt = str(dataset_format).lower().strip()
    validator = FORMAT_VALIDATORS.get(fmt)
    if not validator:
        return [f"不支持的数据集格式: {dataset_format}，支持的格式: {', '.join(sorted(FORMAT_VALIDATORS))}"]
    records = _load_sample(path=path, content=content, max_samples=max_samples)
    return validator(records)