YG_FT/compute/tests/test_llama_factory_adapter.py

from __future__ import annotations

import json

from compute.engines.llama_factory.adapter import _validate_dataset_columns, build_command


def test_build_command_uses_explicit_validation_dataset_without_resplitting() -> None:
    result = build_command(
        {
            "base_model": "/models/qwen",
            "dataset": "ygft_dataset_train",
            "eval_dataset": "ygft_dataset_validation",
            "dataset_dir": "/datasets/example",
            "output_dir": "/outputs/example",
            "val_size": 0.1,
        }
    )

    assert result.command[result.command.index("--dataset") + 1] == "ygft_dataset_train"
    assert result.command[result.command.index("--eval_dataset") + 1] == (
        "ygft_dataset_validation"
    )
    assert "--do_eval" in result.command
    assert "--val_size" not in result.command


def _write(tmp_path, name: str, lines: list[dict]) -> object:
    path = tmp_path / name
    path.write_text(
        "".join(json.dumps(line, ensure_ascii=False) + "\n" for line in lines),
        encoding="utf-8",
    )
    return path


def test_jsonl_alpaca_without_input_column_passes_validation(tmp_path) -> None:
    """纯 jsonl Alpaca 数据缺省 input 字段（常见），不应被校验拦截。"""
    _write(tmp_path, "train.jsonl", [{"instruction": "hi", "output": "hello"}])
    errors = _validate_dataset_columns(
        {
            "dataset_dir": str(tmp_path),
            "dataset_info": {
                "ygft_a": {
                    "file_name": "train.jsonl",
                    "formatting": "alpaca",
                    "columns": {"prompt": "instruction", "query": "input", "response": "output"},
                }
            },
        }
    )
    assert errors == []


def test_jsonl_sharegpt_passes_validation(tmp_path) -> None:
    """ShareGPT 格式 jsonl（messages）应通过校验。"""
    _write(
        tmp_path,
        "msg.jsonl",
        [{"messages": [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}]}],
    )
    errors = _validate_dataset_columns(
        {
            "dataset_dir": str(tmp_path),
            "dataset_info": {
                "ygft_m": {
                    "file_name": "msg.jsonl",
                    "formatting": "sharegpt",
                    "columns": {"messages": "messages"},
                }
            },
        }
    )
    assert errors == []


def test_jsonl_missing_response_still_rejected(tmp_path) -> None:
    """缺 output（response）仍应报错——没有答案无法做有监督微调。"""
    _write(tmp_path, "train.jsonl", [{"instruction": "hi"}])
    errors = _validate_dataset_columns(
        {
            "dataset_dir": str(tmp_path),
            "dataset_info": {
                "ygft_a": {
                    "file_name": "train.jsonl",
                    "formatting": "alpaca",
                    "columns": {"prompt": "instruction", "query": "input", "response": "output"},
                }
            },
        }
    )
    assert errors and "output" in errors[0]