def main()

in grade_school_math/sample.py [0:0]

• 11 lines of code
• 1 McCabe index (conditional complexity)

def main():
    device = th.device("cuda")
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    model = GPT2LMHeadModel.from_pretrained("model_ckpts")
    model.to(device)
    print("Model Loaded")

    test_examples = get_examples("test")
    qn = test_examples[1]["question"]
    sample_len = 100
    print(qn.strip())
    print(sample(model, qn, tokenizer, device, sample_len))