Spaces:
Running
Running
| import torch | |
| import json | |
| from transformers import BertForMaskedLM, BertTokenizer | |
| filename = 'diffusion_lm/improved-diffusion/anlg_results/diff_roc_mbr.json2' | |
| bert_model = 'bert-base-uncased' | |
| tokenizer = BertTokenizer.from_pretrained(bert_model) | |
| model = BertForMaskedLM.from_pretrained(bert_model).cuda() | |
| full_lst = [] | |
| with open(filename, 'r') as f: | |
| for line in f: | |
| line = json.loads(line) | |
| full_lst.append(line) | |
| for example in full_lst: | |
| sent = example['sample'] | |
| obs1 = example['obs1'] | |
| obs2 = example['obs2'] | |
| if 'UNK' in sent: | |
| sent = obs1 + sent.replace('UNK', tokenizer.mask_token) + obs2 | |
| print(sent) | |
| model_inputs = tokenizer(sent,return_tensors="pt") | |
| model_inputs = {k:v.to(model.device) for k,v in model_inputs.items()} | |
| model_out = model(**model_inputs) | |
| mask_words = model_inputs['input_ids'] == tokenizer.mask_token_id | |
| masked_logits = model_out.logits[mask_words].view(-1, model_out.logits.size(-1)) | |
| if masked_logits.size(0) > 0: | |
| # take argmax from this. | |
| max_cands = torch.max(masked_logits, dim=-1) | |
| indices = max_cands.indices | |
| model_inputs['input_ids'][mask_words] = indices | |
| print(tokenizer.batch_decode(model_inputs['input_ids'].tolist())) | |
| else: | |
| print('NO NEED THIS FIX. ') | |