def main()

in src/run_paraphrase.py [0:0]

• 45 lines of code
• 6 McCabe index (conditional complexity)

def main(args):
    roberta = RobertaModel.from_pretrained(args.load_dir, 
                                           checkpoint_file='model.pt',
                                           dropout=0.0)
    # Load model
    roberta.to('cuda')
    roberta.eval()
    print('Finished loading model.')

    # Read test data
    test_data = read_data(args.dataset, roberta)
    print(f'{args.dataset}: Loaded {len(test_data)} test examples.')

    if args.train_ft:
        all_results = []
        for kshot_seed in KSHOT_SEEDS:
            random.seed(args.rng_seed)
            torch.manual_seed(args.rng_seed)
            train_data = read_data(args.dataset, roberta, train=True, kshot_seed=kshot_seed)
            print(f'Loaded {len(train_data)} training examples for kshot-seed {kshot_seed}.')
            dev_data = read_data(args.dataset, roberta, kshot_seed=kshot_seed)
            print(f'Loaded {len(dev_data)} dev examples.')
            model = train_ft(roberta, train_data, dev_data, args.logit_scale, 
                             args.batch_size, args.learning_rate, args.num_epochs)
            pred_scores = predict_ft(test_data, model, roberta.device, args.batch_size)
            stats = evaluate(test_data, pred_scores, threshold=0.0)
            all_results.append(stats)
            print(f'kshot-seed {kshot_seed}: {json.dumps(stats)}')
        print(f'Average: {json.dumps(average_results(all_results))}')
    elif args.train_skl:
        test_features = prune_features(featurize_dataset(test_data, roberta, args.batch_size))
        all_results = []
        for kshot_seed in KSHOT_SEEDS:
            train_data = read_data(args.dataset, roberta, train=True, kshot_seed=kshot_seed)
            print(f'Loaded {len(train_data)} training examples for kshot-seed {kshot_seed}.')
            train_features = prune_features(featurize_dataset(
                    train_data, roberta, args.batch_size))
            vectorizer, scaler, model = train_skl(train_data, train_features, skl_C=args.skl_C)
            pred_scores = predict_skl(test_features, vectorizer, scaler, model)
            stats = evaluate(test_data, pred_scores, threshold=0.0)
            all_results.append(stats)
            print(f'kshot-seed {kshot_seed}: {json.dumps(stats)}')
        print(f'Average: {json.dumps(average_results(all_results))}')
    else:
        feat_dicts = featurize_dataset(test_data, roberta, args.batch_size)
        preds_per_layer = predict_zeroshot(test_data, feat_dicts)
        for name, pred_scores in preds_per_layer.items():
            stats = evaluate(test_data, pred_scores)
            print(f'{name}: {json.dumps(stats)}')