import argparse
from model_loader import load_model
from generator import generate_code
from utils import pretty_print_code

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run Sentinel for code generation")
    parser.add_argument("prompt", type=str, help="The coding prompt to generate code for")
    parser.add_argument("--max_new_tokens", type=int, default=200, help="Maximum tokens to generate")
    args = parser.parse_args()

    generator = load_model("your-username/sentinel")
    code = generate_code(args.prompt, generator, max_new_tokens=args.max_new_tokens)
    pretty_print_code(code)