YG_FT/compute/engines/llama_factory/adapter.py

from __future__ import annotations

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any


@dataclass(frozen=True)
class LlamaFactoryCommand:
    command: list[str]
    work_dir: str
    env: dict[str, str]


def validate_config(config: dict[str, Any]) -> list[str]:
    errors: list[str] = []
    if not config.get("base_model") and not config.get("model_name_or_path"):
        errors.append("base_model or model_name_or_path is required")
    if not config.get("dataset") and not config.get("dataset_dir"):
        errors.append("dataset or dataset_dir is required")
    learning_rate = float(config.get("learning_rate", 0.0002))
    if learning_rate <= 0:
        errors.append("learning_rate must be greater than zero")
    epochs = int(config.get("n_epochs", config.get("num_train_epochs", 1)))
    if epochs <= 0:
        errors.append("n_epochs must be greater than zero")
    return errors


def build_command(config: dict[str, Any], llama_factory_home: str = "/app/LLaMA-Factory") -> LlamaFactoryCommand:
    errors = validate_config(config)
    if errors:
        raise ValueError("; ".join(errors))

    model_path = config.get("base_model") or config.get("model_name_or_path")
    dataset = config.get("dataset") or config.get("dataset_dir")
    output_dir = config.get("output_dir") or f"/data/yg-ft/outputs/{config.get('name', 'training-job')}"
    command = [
        "llamafactory-cli",
        "train",
        "--stage",
        str(config.get("stage", "sft")).lower(),
        "--do_train",
        "true",
        "--model_name_or_path",
        str(model_path),
        "--dataset",
        str(dataset),
        "--template",
        str(config.get("template", "qwen")),
        "--finetuning_type",
        str(config.get("train_method", config.get("finetuning_type", "lora"))),
        "--output_dir",
        str(output_dir),
        "--per_device_train_batch_size",
        str(config.get("batch_size", 2)),
        "--learning_rate",
        str(config.get("learning_rate", 0.0002)),
        "--num_train_epochs",
        str(config.get("n_epochs", 3)),
        "--save_steps",
        str(config.get("save_steps", 50)),
    ]
    quantization_bit = int(config.get("quantization_bit", 0) or 0)
    if quantization_bit in {4, 8}:
        command.extend(["--quantization_bit", str(quantization_bit)])
    return LlamaFactoryCommand(command=command, work_dir=str(Path(llama_factory_home)), env={})


def parse_log_line(line: str) -> dict[str, float] | None:
    if "loss" not in line or "learning_rate" not in line:
        return None
    result: dict[str, float] = {}
    for key in ["loss", "grad_norm", "learning_rate", "epoch"]:
        match = re.search(rf"['\"]?{key}['\"]?\s*:\s*([-+]?\d+(?:\.\d+)?(?:[eE][-+]?\d+)?)", line)
        if match:
            result[key] = float(match.group(1))
    return result or None