Sentinel / inference.py
prelington's picture
Update inference.py
6a5fe57 verified
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)