Views
No views yet
1from transformers import AutoTokenizer, AutoModelForTokenClassification
2import numpy as np
3
4# match tag
5id2tag = {0:'O', 1:'B_MT', 2:'I_MT'}
6
7# load model & tokenizer
8MODEL_NAME = 'MDDDDR/roberta_large_NER'
9
10model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)
11tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
12
13# prepare input
14text = 'mental disorder can also contribute to the development of diabetes through various mechanism including increased stress, poor self care behavior, and adverse effect on glucose metabolism.'
15tokenized = tokenizer(text, return_tensors='pt')
16
17# forward pass
18output = model(**tokenized)
19
20# result
21pred = np.argmax(output[0].cpu().detach().numpy(), axis=2)[0][1:-1]
22
23# check pred
24for txt, pred in zip(tokenizer.tokenize(text), pred):
25 print("{}\t{}".format(id2tag[pred], txt))
26 # B_MT ▁mental
27 # B_MT ▁disorder