Spaces:
Runtime error
Runtime error
File size: 3,092 Bytes
c4ebaf8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from argparse import ArgumentParser
import pickle
import os
import math
import sacrebleu
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelWithLMHead, pipeline, set_seed, GPT2Tokenizer, GPT2Model, MarianTokenizer, MarianMTModel
from constants import *
from model import Model
from util import save_checkpoint, ProgressMeter, AverageMeter, num_params
def avg_formality(preds, model, tokenizer, device='cuda'):
probs = []
for sent in preds:
encoded_input = tokenizer.encode(sent, return_tensors='pt').to(device)
lengths = torch.LongTensor([encoded_input.shape[1]]).to(device)
scores = model(encoded_input, lengths=lengths) # batch x seq
score = scores.flatten()[-1].item()
probs.append(math.exp(score) / (1 + math.exp(score))) # sigmoided score = prob
return np.mean(probs)
if __name__=='__main__':
parser = ArgumentParser()
parser.add_argument('--pred', type=str)
parser.add_argument('--ref', type=str, nargs='*', help='bleu refs')
parser.add_argument('--ckpt', type=str, help='formality classifier')
parser.add_argument('--dataset_info', type=str)
parser.add_argument('--device', type=str, default='cuda', choices=['cpu', 'cuda'])
parser.add_argument('--model_string', type=str, default='Helsinki-NLP/opus-mt-es-en')
args = parser.parse_args()
# refs = [['The dog bit the man.', 'It was not unexpected.', 'The man bit him first.'],
# ['The dog had bit the man.', 'No one was surprised.', 'The man had bitten the dog.']]
# sys = ['The dog bit the man.', "It wasn't surprising.", 'The man had just bitten him.']
print('num ref files', len(args.ref))
pred = []
with open(args.pred, 'r') as rf:
for line in rf:
pred.append(line.strip())
refs = []
for ref_file in args.ref:
ref = []
with open(ref_file, 'r') as rf:
for line in rf:
ref.append(line.strip())
assert len(ref) == len(pred)
refs.append(ref)
bleu = sacrebleu.corpus_bleu(pred, refs)
print('BLEU score:', bleu.score)
with open(args.dataset_info, 'rb') as rf:
dataset_info = pickle.load(rf)
tokenizer = MarianTokenizer.from_pretrained(args.model_string)
tokenizer.add_special_tokens({'pad_token': PAD_TOKEN})
pad_id = tokenizer.encode(PAD_TOKEN)[0]
checkpoint = torch.load(args.ckpt, map_location=args.device)
model_args = checkpoint['args']
conditioning_model = Model(model_args, pad_id, len(dataset_info.index2word)) # no need to get the glove embeddings when reloading since they're saved in model ckpt anyway
conditioning_model.load_state_dict(checkpoint['state_dict'])
conditioning_model = conditioning_model.to(args.device)
conditioning_model.eval()
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.ckpt, checkpoint['epoch']))
print('num params', num_params(conditioning_model))
print('avg formality prob according to model', avg_formality(pred, conditioning_model, tokenizer, device=args.device))
|