finetune/scripts/eval_actor_gguf.py

#!/usr/bin/env python3
"""Evaluate an Actor GGUF model through llama.cpp prompt formats."""

from __future__ import annotations

import argparse
import json
import re
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Any

from eval_minicpm5_actor_lora import load_eval_rows, validate_generation


DEFAULT_MODEL = Path("finetune/outputs/gguf/minicpm5-actor-q4_k_m.gguf")
DEFAULT_LLAMA_BIN = Path("../llama.cpp/build/bin/llama-completion")
DEFAULT_EVAL_FILE = Path("finetune/data_samples/actor_eval_prompts.jsonl")
DEFAULT_OUTPUT_FILE = Path("finetune/eval_outputs/minicpm5_actor_gguf_eval.jsonl")
PROMPT_FORMATS = ["raw", "simple_json", "chatml", "system_user_assistant"]
STOP_STRINGS = [
    "\nUSER:",
    "\nSYSTEM:",
    "\nASSISTANT:",
    "\nJSON:",
    "[Start thinking]",
    "</s>",
]
LOG_LINE_RE = re.compile(
    r"^\s*(?:\d+\.\d+\.\d+\s+[A-Z]\s+)?(?:llama_|ggml_|common_|sampling_|sampler|main:|system_info:|perf:)"
)
METRIC_KEYS = [
    "raw_clean_json",
    "extracted_json_parse",
    "missing_required_fields_success",
    "exact_top_level_schema_success",
    "extra_top_level_fields",
    "forbidden_top_level_fields",
    "sanitized_actor_json_usable",
    "strict_tool_request",
    "sanitized_tool_request_usable",
    "line_length_pass",
    "interactive_marker_seen",
    "runtime_error",
    "timeout",
]
METRIC_LABELS = {
    "raw_clean_json": "raw clean JSON",
    "extracted_json_parse": "extracted JSON parse",
    "missing_required_fields_success": "has required fields",
    "exact_top_level_schema_success": "exact top-level schema",
    "extra_top_level_fields": "extra top-level fields",
    "forbidden_top_level_fields": "forbidden top-level fields",
    "sanitized_actor_json_usable": "sanitized actor JSON usable",
    "strict_tool_request": "strict tool_request",
    "sanitized_tool_request_usable": "sanitized tool_request usable",
    "line_length_pass": "line length pass",
    "interactive_marker_seen": "interactive marker seen",
    "runtime_error": "runtime error",
    "timeout": "timeout",
}
SCHEMA_REMINDER = (
    "Return exactly one JSON object with exactly these keys: "
    "intent, line, emotion, gesture, stage_effect, memory_update, tool_request. "
    "Do not omit stage_effect. Do not include markdown, commentary, copied input fields, or another assistant turn. "
    "Stop after the JSON object."
)
JSON_SCHEMA_EXAMPLE = (
    '{"intent":"react_to_event","line":"A short puppet line tied to the scene.",'
    '"emotion":"curious","gesture":"tilts head toward the stage lights",'
    '"stage_effect":"spotlight_glow","memory_update":null,"tool_request":null}'
)
INTERACTIVE_MARKERS = [
    "interactive mode on",
    "available commands:",
    "chat template is available, enabling conversation mode",
    "\n>",
    "please use llama-completion instead",
    "[Start thinking]",
]


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", type=Path, default=DEFAULT_MODEL)
    parser.add_argument(
        "--llama_bin",
        "--llama_cli",
        dest="llama_bin",
        type=Path,
        default=DEFAULT_LLAMA_BIN,
        help="Path to llama.cpp completion binary. Defaults to llama-completion; --llama_cli is a backward-compatible alias.",
    )
    parser.add_argument("--eval_file", type=Path, default=DEFAULT_EVAL_FILE)
    parser.add_argument("--output_file", type=Path, default=DEFAULT_OUTPUT_FILE)
    parser.add_argument("--limit", type=int, default=None)
    parser.add_argument("--n_predict", type=int, default=160)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--top_p", type=float, default=0.9)
    parser.add_argument(
        "--prompt_format",
        choices=["auto", *PROMPT_FORMATS],
        default="raw",
        help="Use one prompt format, or auto to run every supported format.",
    )
    parser.add_argument(
        "--timeout_seconds",
        type=int,
        default=120,
        help="Per-prompt llama.cpp timeout.",
    )
    parser.add_argument(
        "--no_stop_strings",
        action="store_true",
        help="Skip best-effort llama.cpp reverse-prompt stop strings.",
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> None:
    args = parse_args(argv)
    run_eval(args)


def run_eval(args: argparse.Namespace) -> None:
    model_path = args.model.expanduser()
    if not model_path.exists():
        raise SystemExit(
            f"GGUF model not found: {model_path}\n"
            "Create it first with finetune/scripts/convert_actor_merged_to_gguf.sh."
        )
    if not model_path.is_file():
        raise SystemExit(f"GGUF model path is not a file: {model_path}")
    llama_bin = resolve_llama_bin(args.llama_bin)
    if not args.eval_file.exists():
        raise SystemExit(f"Eval prompt file not found: {args.eval_file}")
    binary_help = read_binary_help(llama_bin)
    binary_capabilities = detect_binary_capabilities(binary_help)
    if llama_bin.name == "llama-cli":
        print(
            "warning: llama-cli may enter interactive mode for this llama.cpp build; "
            "prefer llama-completion for one-shot GGUF eval.",
            file=sys.stderr,
        )

    rows = load_eval_rows(args.eval_file, args.limit)
    formats = PROMPT_FORMATS if args.prompt_format == "auto" else [args.prompt_format]
    args.output_file.parent.mkdir(parents=True, exist_ok=True)

    counts_by_format = {prompt_format: new_metric_counts() for prompt_format in formats}
    totals_by_format = {prompt_format: 0 for prompt_format in formats}

    with args.output_file.open("w", encoding="utf-8") as handle:
        for row_index, row in enumerate(rows, start=1):
            for prompt_format in formats:
                row_id = row.get("id", f"eval-{row_index:03d}")
                print(f"evaluating prompt_format={prompt_format} row_id={row_id}", flush=True)
                prompt = build_prompt(row["messages"], prompt_format)
                raw_output, command_info = run_llama_completion(
                    llama_bin=llama_bin,
                    model_path=model_path,
                    prompt=prompt,
                    n_predict=args.n_predict,
                    temperature=args.temperature,
                    top_p=args.top_p,
                    use_stop_strings=not args.no_stop_strings,
                    timeout_seconds=args.timeout_seconds,
                    capabilities=binary_capabilities,
                )
                validation = validate_generation(raw_output)
                validation = apply_gguf_empty_tool_list_sanitization(validation)
                validation["start_thinking_seen"] = "[Start thinking]" in raw_output
                validation["interactive_marker_seen"] = has_interactive_marker(raw_output) or has_interactive_marker(command_info["stderr"])
                validation["runtime_error"] = bool(command_info["runtime_error"])
                validation["timeout"] = bool(command_info["timeout"])
                for metric in METRIC_KEYS:
                    counts_by_format[prompt_format][metric] += int(validation[metric])
                counts_by_format[prompt_format]["start_thinking_seen"] += int(validation["start_thinking_seen"])
                if validation["interactive_marker_seen"]:
                    print(
                        f"warning: interactive marker seen for row {row.get('id', row_index)} "
                        f"with prompt_format={prompt_format}",
                        file=sys.stderr,
                    )
                if validation["runtime_error"]:
                    print(
                        f"warning: llama.cpp runtime error for row {row_id} "
                        f"with prompt_format={prompt_format}: {command_info['runtime_error']}",
                        file=sys.stderr,
                        flush=True,
                    )
                totals_by_format[prompt_format] += 1

                detail = {
                    "id": row_id,
                    "row_type": row.get("row_type"),
                    "prompt_format": prompt_format,
                    "prompt": prompt,
                    "raw_output": raw_output,
                    "raw_stdout": command_info["stdout"],
                    "raw_stderr": command_info["stderr"],
                    "json_text": validation["json_text"],
                    "parsed_json": validation["parsed_json"],
                    "sanitized_actor_json": validation["sanitized_actor_json"],
                    "validation": {
                        key: value
                        for key, value in validation.items()
                        if key
                        not in {
                            "parsed_json",
                            "sanitized_actor_json",
                            "json_text",
                            "original_tool_request",
                            "sanitized_tool_request",
                        }
                    },
                    "original_tool_request": validation["original_tool_request"],
                    "sanitized_tool_request": validation["sanitized_tool_request"],
                    "llama_cpp": command_info,
                }
                handle.write(json.dumps(detail, ensure_ascii=True, separators=(",", ":")) + "\n")

    for prompt_format in formats:
        print_format_summary(prompt_format, totals_by_format[prompt_format], counts_by_format[prompt_format])
    print(f"wrote detailed generations to: {args.output_file}")


def resolve_llama_bin(raw_path: Path) -> Path:
    expanded = raw_path.expanduser()
    if expanded.exists():
        if expanded.is_file():
            return expanded
        raise SystemExit(f"llama.cpp binary path is not a file: {expanded}")
    found = shutil.which(str(raw_path))
    if found:
        return Path(found)
    raise SystemExit(
        f"llama.cpp binary not found: {raw_path}\n"
        "Pass --llama_bin /absolute/path/to/llama.cpp/build/bin/llama-completion."
    )


def build_prompt(messages: list[dict[str, str]], prompt_format: str) -> str:
    system, user = split_system_user(messages)
    if prompt_format == "raw":
        return f"{system.strip()}\n\n{user.strip()}\n\n{SCHEMA_REMINDER}\nRequired shape:\n{JSON_SCHEMA_EXAMPLE}\n\nAssistant JSON:\n"
    if prompt_format == "simple_json":
        return (
            f"System: {system.strip()}\n\n"
            f"User:\n{user.strip()}\n\n"
            f"Output contract: {SCHEMA_REMINDER}\n\n"
            f"Required JSON shape:\n{JSON_SCHEMA_EXAMPLE}\n\n"
            "JSON:\n"
        )
    if prompt_format == "chatml":
        return (
            f"<|im_start|>system\n{system.strip()}\n<|im_end|>\n"
            f"<|im_start|>user\n{user.strip()}\n\n{SCHEMA_REMINDER}\nRequired JSON shape:\n{JSON_SCHEMA_EXAMPLE}\n<|im_end|>\n"
            "<|im_start|>assistant\n"
        )
    if prompt_format == "system_user_assistant":
        return (
            f"### SYSTEM\n{system.strip()}\n\n"
            f"### USER\n{user.strip()}\n\n{SCHEMA_REMINDER}\nRequired JSON shape:\n{JSON_SCHEMA_EXAMPLE}\n\n"
            "### ASSISTANT\n"
        )
    raise ValueError(f"Unsupported prompt format: {prompt_format}")


def split_system_user(messages: list[dict[str, str]]) -> tuple[str, str]:
    system_parts = [message["content"] for message in messages if message.get("role") == "system"]
    user_parts = [message["content"] for message in messages if message.get("role") == "user"]
    system = "\n\n".join(system_parts)
    user = "\n\n".join(user_parts)
    if not system:
        system = "You are an Actor agent in AI Puppet Theater. Return only one valid JSON object."
    if not user:
        raise ValueError("Eval row is missing a user message")
    return system, user


def run_llama_completion(
    *,
    llama_bin: Path,
    model_path: Path,
    prompt: str,
    n_predict: int,
    temperature: float,
    top_p: float,
    use_stop_strings: bool,
    timeout_seconds: int,
    capabilities: dict[str, bool],
) -> tuple[str, dict[str, Any]]:
    base_command = [
        str(llama_bin),
        "-m",
        str(model_path),
        "-p",
        prompt,
        "-n",
        str(n_predict),
        "--temp",
        str(temperature),
        "--top-p",
        str(top_p),
    ]
    optional_args: list[str] = []
    if capabilities["no_conversation_short"]:
        optional_args.append("-no-cnv")
    elif capabilities["no_conversation_long"]:
        optional_args.append("--no-conversation")
    if capabilities["no_display_prompt"]:
        optional_args.append("--no-display-prompt")
    if capabilities["reasoning"]:
        optional_args.extend(["--reasoning", "off"])
    if capabilities["reasoning_budget"]:
        optional_args.extend(["--reasoning-budget", "0"])
    if use_stop_strings and capabilities["reverse_prompt"]:
        optional_args.extend(["--reverse-prompt", ",".join(STOP_STRINGS)])

    result = run_command(base_command + optional_args, timeout_seconds)
    used_optional_args = True
    if result["returncode"] != 0 and looks_like_cli_option_error(result["stderr"]):
        result = run_command(base_command, timeout_seconds)
        used_optional_args = False

    runtime_error = None
    if result["timeout"]:
        runtime_error = f"timeout_after_{timeout_seconds}_seconds"
    elif result["returncode"] != 0:
        runtime_error = f"nonzero_exit_{result['returncode']}"

    output = extract_candidate_generation(result["stdout"], prompt)
    return output.strip(), {
        "path": str(llama_bin),
        "command": base_command + (optional_args if used_optional_args else []),
        "used_optional_args": used_optional_args,
        "capabilities": capabilities,
        "returncode": result["returncode"],
        "timeout": result["timeout"],
        "runtime_error": runtime_error,
        "stdout": result["stdout"].strip(),
        "stderr": result["stderr"].strip(),
    }


def read_binary_help(llama_bin: Path) -> str:
    result = subprocess.run(
        [str(llama_bin), "--help"],
        check=False,
        capture_output=True,
        text=True,
        timeout=20,
    )
    return f"{result.stdout}\n{result.stderr}"


def detect_binary_capabilities(help_text: str) -> dict[str, bool]:
    return {
        "no_conversation_short": "-no-cnv" in help_text,
        "no_conversation_long": "--no-conversation" in help_text,
        "no_display_prompt": "--no-display-prompt" in help_text,
        "reasoning": "--reasoning" in help_text,
        "reasoning_budget": "--reasoning-budget" in help_text,
        "reverse_prompt": "--reverse-prompt" in help_text or "-r," in help_text or "-r " in help_text,
    }


def run_command(command: list[str], timeout_seconds: int) -> dict[str, Any]:
    try:
        result = subprocess.run(
            command,
            check=False,
            capture_output=True,
            text=True,
            timeout=timeout_seconds,
        )
    except subprocess.TimeoutExpired as exc:
        return {
            "returncode": 124,
            "stdout": exc.stdout or "",
            "stderr": exc.stderr or f"timed out after {timeout_seconds} seconds",
            "timeout": True,
        }
    return {
        "returncode": result.returncode,
        "stdout": result.stdout,
        "stderr": result.stderr,
        "timeout": False,
    }


def looks_like_cli_option_error(stderr: str) -> bool:
    lowered = stderr.lower()
    option_markers = [
        "unknown argument",
        "unknown option",
        "unrecognized option",
        "invalid argument",
    ]
    optional_flag_markers = ["no-cnv", "no-conversation", "no-display-prompt", "reverse-prompt", "reasoning", "reasoning-budget"]
    return any(marker in lowered for marker in option_markers) and any(marker in lowered for marker in optional_flag_markers)


def extract_candidate_generation(stdout: str, prompt: str) -> str:
    output = stdout.replace("\r\n", "\n")
    if output.startswith(prompt):
        output = output[len(prompt) :]
    lines = []
    for line in output.splitlines():
        if LOG_LINE_RE.search(line):
            continue
        lines.append(line)
    cleaned = "\n".join(lines).strip()
    json_start = cleaned.find("{")
    if json_start > 0:
        prefix = cleaned[:json_start].strip()
        if prefix.upper() in {"JSON:", "ASSISTANT JSON:", "ASSISTANT:"} or prefix.endswith("JSON:"):
            cleaned = cleaned[json_start:]
    return cleaned


def apply_gguf_empty_tool_list_sanitization(validation: dict[str, Any]) -> dict[str, Any]:
    parsed_json = validation.get("parsed_json")
    if not isinstance(parsed_json, dict) or parsed_json.get("tool_request") != []:
        return validation

    sanitized_actor_json = validation.get("sanitized_actor_json")
    if isinstance(sanitized_actor_json, dict):
        sanitized_actor_json = dict(sanitized_actor_json)
        sanitized_actor_json["tool_request"] = None
        validation["sanitized_actor_json"] = sanitized_actor_json
        if validation["missing_required_fields_success"]:
            validation["sanitized_actor_json_usable"] = True

    validation["sanitized_tool_request"] = None
    validation["sanitized_tool_request_usable"] = True
    validation["strict_tool_request"] = False
    validation["sanitization_needed"] = True
    validation["actor_sanitization_needed"] = True
    validation["tool_error"] = "tool_request_empty_list_sanitized_to_null"
    return validation


def new_metric_counts() -> dict[str, int]:
    counts = {key: 0 for key in METRIC_KEYS}
    counts["start_thinking_seen"] = 0
    return counts


def print_format_summary(prompt_format: str, total: int, metric_counts: dict[str, int]) -> None:
    print(f"prompt_format: {prompt_format}")
    print(f"total generations: {total}")
    for key in METRIC_KEYS:
        count = metric_counts[key]
        rate = count / total if total else 0.0
        print(f"{METRIC_LABELS[key]}: {count}/{total} ({rate:.1%})")
    start_count = metric_counts["start_thinking_seen"]
    start_rate = start_count / total if total else 0.0
    print(f"[Start thinking] seen: {start_count}/{total} ({start_rate:.1%})")


def has_interactive_marker(text: str) -> bool:
    lowered = text.lower()
    return any(marker in lowered for marker in INTERACTIVE_MARKERS)


def format_command(command: list[str]) -> str:
    return " ".join(shlex_quote(part) for part in command)


def shlex_quote(value: str) -> str:
    if not value:
        return "''"
    safe_chars = set("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_+-=.,/:@%")
    if all(char in safe_chars for char in value):
        return value
    return "'" + value.replace("'", "'\"'\"'") + "'"


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        sys.exit(130)