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/dmis_lab_biobert_v1.1_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
21preds = np.argmax(output[0].cpu().detach().numpy(), axis=2)[0][1:-1]
22
23# check preds
24for txt, pred in zip(tokenizer.tokenize(text), preds):
25 print("{}\t{}".format(id2tag[pred], txt))
26 # B_MT mental
27 # B_MT disorder
28 # O can
29 # O also
30 # O contribute
31 # O to
32 # O the
33 # B_MT development
34 # O of
35 # B_MT diabetes
36 # O through
37 # O various
38 # B_MT mechanism
39 # O including
40 # O increased
41 # B_MT stress
42 # O ,
43 # O poor
44 # B_MT self
45 # B_MT care
46 # B_MT behavior
47 # O ,
48 # O and
49 # B_MT adverse
50 # I_MT effect
51 # O on
52 # B_MT glucose
53 # B_MT metabolism
54 # O .