generate.py

#!/usr/bin/env python3
"""Run the pinned Guide Labs Steerling-8B quickstart as a finite GPU job."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

MODEL_ID = "guidelabs/steerling-8b"
MODEL_REVISION = "f2e59522c2f6d91c063770a4ce5514947a7b68bc"
SOURCE_REVISION = "f34ffa89e46969445f3cf6e7c885e9623a2047c1"
MINIMUM_VRAM_BYTES = 18 * 1024**3


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate text with the pinned Guide Labs Steerling-8B model."
    )
    parser.add_argument(
        "--prompt",
        default="The key to understanding neural networks is",
    )
    parser.add_argument("--max-new-tokens", type=int, default=128)
    parser.add_argument("--steps", type=int, default=128)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--output",
        default="/mnt/output/steerling.json",
        help="JSON result path; /mnt/output is persisted when an output bucket is mounted.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    import torch
    from steerling import GenerationConfig, SteerlingGenerator
    from transformers import AutoModel, AutoTokenizer

    if not torch.cuda.is_available():
        raise RuntimeError("Steerling-8B requires a CUDA GPU with at least 18 GiB VRAM")

    device = torch.cuda.current_device()
    properties = torch.cuda.get_device_properties(device)
    if properties.total_memory < MINIMUM_VRAM_BYTES:
        available_gib = properties.total_memory / 1024**3
        raise RuntimeError(
            f"Steerling-8B requires at least 18 GiB VRAM; found {available_gib:.1f} GiB"
        )

    model = AutoModel.from_pretrained(
        MODEL_ID,
        revision=MODEL_REVISION,
        trust_remote_code=True,
        dtype=torch.bfloat16,
        device_map="cuda",
        low_cpu_mem_usage=True,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        revision=MODEL_REVISION,
        trust_remote_code=True,
    )
    generator = SteerlingGenerator.from_model(model, tokenizer, device="cuda")
    text = generator.generate(
        args.prompt,
        GenerationConfig(
            max_new_tokens=args.max_new_tokens,
            steps=args.steps,
            seed=args.seed,
        ),
    )

    result = {
        "model": MODEL_ID,
        "model_revision": MODEL_REVISION,
        "source_revision": SOURCE_REVISION,
        "gpu": properties.name,
        "prompt": args.prompt,
        "max_new_tokens": args.max_new_tokens,
        "steps": args.steps,
        "seed": args.seed,
        "text": text,
    }
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()