approach0/dpr-cotmae-120

1import re2import os3import fire4import torch5from functools import partial6from transformers import AutoTokenizer7from transformers import AutoModelForPreTraining8from pya0.preprocess import preprocess_for_transformer9 10 11def highlight_masked(txt):12    return re.sub(r"(\[MASK\])", '\033[92m' + r"\1" + '\033[0m', txt)13 14 15def classifier_hook(tokenizer, tokens, topk, module, inputs, outputs):16    unmask_scores, seq_rel_scores = outputs17    MSK_CODE = 10318    token_ids = tokens['input_ids'][0]19    masked_idx = (token_ids == torch.tensor([MSK_CODE]))20    scores = unmask_scores[0][masked_idx]21    cands = torch.argsort(scores, dim=1, descending=True)22    for i, mask_cands in enumerate(cands):23        top_cands = mask_cands[:topk].detach().cpu()24        print(f'MASK[{i}] top candidates: ' +25            str(tokenizer.convert_ids_to_tokens(top_cands)))26 27 28def test(tokenizer_name_or_path, model_name_or_path, test_file='test.txt'):29 30    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)31    model = AutoModelForPreTraining.from_pretrained(model_name_or_path,32        tie_word_embeddings=True33    )34    with open(test_file, 'r') as fh:35        for line in fh:36            # parse test file line37            line = line.rstrip()38            fields = line.split('\t')39            maskpos = list(map(int, fields[0].split(',')))40            # preprocess and mask words41            sentence = preprocess_for_transformer(fields[1])42            tokens = sentence.split()43            for pos in filter(lambda x: x!=0, maskpos):44                tokens[pos-1] = '[MASK]'45            sentence = ' '.join(tokens)46            sentence = sentence.replace('[mask]', '[MASK]')47            tokens = tokenizer(sentence,48                padding=True, truncation=True, return_tensors="pt")49            #print(tokenizer.decode(tokens['input_ids'][0]))50            print('*', highlight_masked(sentence))51            # print unmasked52            with torch.no_grad():53                display = ['\n', '']54                classifier = model.cls55                partial_hook = partial(classifier_hook, tokenizer, tokens, 3)56                hook = classifier.register_forward_hook(partial_hook)57                model(**tokens)58                hook.remove()59 60 61if __name__ == '__main__':62    os.environ["PAGER"] = 'cat'63    fire.Fire(test)64