scripts/test_lora_model.py

#!/usr/bin/env python3
"""Load base model + LoRA adapter and run regression eval prompts."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))

from mythos_lora_core import DEFAULT_MAX_NEW_TOKENS, run_eval


def project_root() -> Path:
    return Path(__file__).resolve().parent.parent


def parse_args() -> argparse.Namespace:
    root = project_root()
    parser = argparse.ArgumentParser(description="Test Mythos-Coder LoRA adapter.")
    parser.add_argument(
        "--model_name",
        default="Qwen/Qwen2.5-Coder-0.5B-Instruct",
        help="Hugging Face base model ID",
    )
    parser.add_argument(
        "--adapter_dir",
        default=str(root / "models" / "mythos-coder-qwen-0.5b-lora"),
        help="LoRA adapter directory",
    )
    parser.add_argument(
        "--prompts_file",
        default=str(root / "data" / "eval" / "code_output_regression_prompts.jsonl"),
        help="Eval prompts JSONL",
    )
    parser.add_argument(
        "--output_file",
        default=str(root / "data" / "eval" / "runpod_lora_results.jsonl"),
        help="Results JSONL output path",
    )
    parser.add_argument(
        "--max_new_tokens",
        type=int,
        default=DEFAULT_MAX_NEW_TOKENS,
        help="Max tokens to generate per prompt",
    )
    return parser.parse_args()


def resolve_path(path_str: str) -> Path:
    path = Path(path_str)
    if not path.is_absolute():
        path = project_root() / path
    return path


def main() -> int:
    args = parse_args()
    prompts_path = resolve_path(args.prompts_file)
    model_path = resolve_path(args.adapter_dir)
    output_path = resolve_path(args.output_file)

    if not prompts_path.exists():
        print(f"Error: prompts file not found: {prompts_path}", file=sys.stderr)
        return 1
    if not model_path.exists():
        print(f"Error: adapter directory not found: {model_path}", file=sys.stderr)
        return 1

    print("=== Mythos-Coder LoRA Eval ===")
    print(f"CUDA available: {__import__('torch').cuda.is_available()}")

    result = run_eval(
        prompts_path,
        model_path,
        output_path,
        model_name=args.model_name,
        max_new_tokens=args.max_new_tokens,
    )

    print("=== Eval Summary ===")
    print(f"Prompts tested:     {result['prompt_count']}")
    print(f"Results written to: {result['output_path']}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())