| | import os,csv,time
|
| | from argparse import ArgumentParser
|
| |
|
| | import numpy as np
|
| | import torch
|
| | import torch.nn.functional as F
|
| |
|
| | from transformers import AutoModelForCausalLM
|
| | from tqdm import tqdm
|
| |
|
| | from anticipation.config import EVENT_SIZE
|
| |
|
| | def log_loss(model, datafile, subsample):
|
| | with open(datafile, 'r') as data:
|
| | ce = torch.empty(0)
|
| | for i,line in tqdm(list(enumerate(data))):
|
| | if i % subsample != 0:
|
| | continue
|
| |
|
| | tokens = [int(token) for token in line.split()]
|
| | tokens = torch.tensor(tokens).unsqueeze(0).cuda()
|
| | with torch.no_grad():
|
| | logits = model(tokens).logits[0]
|
| | ce = torch.cat([ce, F.cross_entropy(logits[:-1],tokens[0,1:],reduction='none').cpu()])
|
| |
|
| | return ce
|
| |
|
| |
|
| | def main(args):
|
| | print(f'Sub-sampling results at rate {args.subsample}')
|
| |
|
| | results = os.path.join(args.model, args.output)
|
| | print(f'Storing results at {results}')
|
| |
|
| | checkpoints = [os.path.join(f.path, 'hf') for f in os.scandir(args.model) if
|
| | f.is_dir() and os.path.basename(f).startswith('step-')]
|
| |
|
| | if args.all:
|
| | print('Calculating log-loss for checkpoints:')
|
| | for ckpt in checkpoints:
|
| | print(' ', ckpt)
|
| | else:
|
| | steps = [int(ckpt.split(os.sep)[-2][5:]) for ckpt in checkpoints]
|
| | checkpoints = [os.path.join(args.model, f'step-{max(steps)}', 'hf')]
|
| | print('Calculating log-loss for final checkpoint:')
|
| | print(' ', checkpoints[0])
|
| |
|
| | print('Calculating log-loss on dataset:')
|
| | print(' ', args.filename)
|
| | with open(results, 'w', newline='') as f:
|
| | fields = ['step', 'loss']
|
| | if args.bpe:
|
| | fields.append('bpe')
|
| | if not args.interarrival:
|
| | fields.extend(['event_ppl', 'onset_ppl', 'dur_ppl', 'note_ppl'])
|
| |
|
| | writer = csv.DictWriter(f, fieldnames=fields)
|
| | writer.writeheader()
|
| | for ckpt in checkpoints:
|
| | step = int(ckpt.split(os.sep)[-2][5:])
|
| | print(f'Loading checkpoint (step {step}):')
|
| | print(' ', ckpt)
|
| | t0 = time.time()
|
| | model = AutoModelForCausalLM.from_pretrained(ckpt).cuda()
|
| | print(f' loaded in {time.time()-t0} seconds')
|
| |
|
| | ce = log_loss(model, args.filename, args.subsample)
|
| |
|
| | res = {}
|
| | res['step'] = step
|
| | res['loss'] = np.round(ce.mean().item(), 3)
|
| | if args.bpe:
|
| |
|
| | assert os.path.basename(args.filename) == 'test.txt'
|
| | res['bpe'] = args.subsample*ce.mean().item()*np.log2(np.e)*(len(ce) / (560.98*3600))
|
| | if not args.interarrival:
|
| | res['event_ppl'] = np.round(np.exp(EVENT_SIZE*ce.mean().item()), 3)
|
| | res['onset_ppl'] = np.round(np.exp(ce[0::3].mean().item()), 3)
|
| | res['dur_ppl'] = np.round(np.exp(ce[1::3].mean().item()), 3)
|
| | res['note_ppl'] = np.round(np.exp(ce[2::3].mean().item()), 3)
|
| |
|
| | writer.writerow(res)
|
| |
|
| |
|
| | if __name__ == '__main__':
|
| | parser = ArgumentParser(description='evaluate log-loss for a tokenized dataset')
|
| | parser.add_argument('-f', '--filename', help='file containing a tokenized dataset')
|
| | parser.add_argument('-m', '--model', help='file containing a model to evaluate')
|
| | parser.add_argument('-o', '--output', help='output file')
|
| | parser.add_argument('-v', '--verbose', action='store_true', help='verbose console output')
|
| | parser.add_argument('-a', '--all', action='store_true',
|
| | help='calculate loss for all checkpoints')
|
| | parser.add_argument('--bpe', action='store_true',
|
| | help='calculate loss for all checkpoints')
|
| | parser.add_argument('-i', '--interarrival', action='store_true',
|
| | help='request interarrival-time enocoding (default to arrival-time encoding)')
|
| | parser.add_argument('-s', '--subsample', type=int, default=10,
|
| | help='dataset subsampling ratio')
|
| |
|
| | main(parser.parse_args())
|
| |
|
