Update README.md

README.md

@@ -33,7 +33,6 @@ from gec.tag import rewrite
 
 tokenizer = BertTokenizer.from_pretrained('CAMeL-Lab/text-editing-coda')
 model = BertForTokenClassification.from_pretrained('CAMeL-Lab/text-editing-coda')
-edits_map = model.config.id2label
 
 text = 'أنا بعطيك رقم تلفونو و عنوانو'.split()
 
@@ -43,7 +42,7 @@ with torch.no_grad():
     logits = model(**tokenized_text).logits
     preds = F.softmax(logits.squeeze(), dim=-1)
     preds = torch.argmax(preds, dim=-1).cpu().numpy()
-    edits = [edits_map[p] for p in preds[1:-1]]
+    edits = [model.config.id2label[p] for p in preds[1:-1]]
     assert len(edits) == len(tokenized_text['input_ids'][0][1:-1])
 
 print(edits) # ['R_[ا]K*', 'K*I_[ا]K', 'K*', 'K*', 'K*', 'K*', 'K*R_[ه]', 'K*', 'MK*', 'R_[ه]']