scripts/check_writer_premises.py

"""Check that the active writer returns schema-valid scripts for smoke premises."""

from __future__ import annotations

import argparse
import json

from midnight_static.schema import Genre
from midnight_static.writer import SMOKE_PREMISES, NemotronWriter, get_writer, script_to_json_dict


def main() -> None:
    parser = argparse.ArgumentParser(description="Run writer schema smoke checks.")
    parser.add_argument("--writer", default="fixture", help="Writer implementation name.")
    parser.add_argument("--model-id", help="Optional HF model id for --writer nemotron.")
    parser.add_argument("--max-new-tokens", type=int, help="Optional generation cap.")
    parser.add_argument("--limit", type=int, default=len(SMOKE_PREMISES), help="Number of premises.")
    args = parser.parse_args()

    if args.writer == "nemotron" and (args.model_id or args.max_new_tokens):
        writer = NemotronWriter(model_id=args.model_id, max_new_tokens=args.max_new_tokens)
    else:
        writer = get_writer(args.writer)
    results = []
    for premise in SMOKE_PREMISES[: args.limit]:
        script = writer.write(premise, genre=Genre.WEIRD)
        script.validate()
        results.append(
            {
                "premise": premise,
                "writer": writer.name,
                "title": script.title,
                "genre": script.genre.value,
                "cast_count": len(script.cast),
                "scene_count": len(script.scenes),
                "line_count": sum(len(scene.lines) for scene in script.scenes),
                "json": script_to_json_dict(script),
            }
        )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()