Views
No views yet
1from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoConfig
2from transformers import pipeline
3
4id2label = {
5'LABEL_0': 'B-NATIONAL',
6 'LABEL_1': 'I-PLANT',
7 'LABEL_2': 'I-ORGANISATION',
8 'LABEL_3': 'B-ORGANISATION',
9 'LABEL_4': 'B-MEDIA',
10 'LABEL_5': 'I-ARTIFACT',
11 'LABEL_6': 'B-AWARD',
12 'LABEL_7': 'B-UNKNOWN',
13 'LABEL_8': 'I-LOCATION',
14 'LABEL_9': 'B-PERSON',
15 'LABEL_10': 'I-LEGAL',
16 'LABEL_11': 'B-BUSINESS',
17 'LABEL_12': 'B-ACRONYM',
18 'LABEL_13': 'I-PERIOD',
19 'LABEL_14': 'B-INSTITUTION',
20 'LABEL_15': 'I-MEASURE',
21 'LABEL_16': 'B-CREATION',
22 'LABEL_17': 'I-ACRONYM',
23 'LABEL_18': 'I-AWARD',
24 'LABEL_19': 'I-WEBSITE',
25 'LABEL_20': 'B-PERIOD',
26 'LABEL_21': 'I-PERSON',
27 'LABEL_22': 'I-PERSON_TYPE',
28 'LABEL_23': 'B-SUBSTANCE',
29 'LABEL_24': 'O',
30 'LABEL_25': 'B-PLANT',
31 'LABEL_26': 'I-INSTITUTION',
32 'LABEL_27': 'I-SUBSTANCE',
33 'LABEL_28': 'I-INSTALLATION',
34 'LABEL_29': 'B-CONCEPT',
35 'LABEL_30': 'B-TITLE',
36 'LABEL_31': 'I-EVENT',
37 'LABEL_32': 'B-ARTIFACT',
38 'LABEL_33': 'B-MEASURE',
39 'LABEL_34': 'B-LOCATION',
40 'LABEL_35': 'I-BUSINESS',
41 'LABEL_36': 'B-ANIMAL',
42 'LABEL_37': 'B-PERSON_TYPE',
43 'LABEL_38': 'B-INSTALLATION',
44 'LABEL_39': 'I-TITLE',
45 'LABEL_40': 'B-IDENTIFIER',
46 'LABEL_41': 'I-IDENTIFIER',
47 'LABEL_42': 'B-LEGAL',
48 'LABEL_43': 'I-MEDIA',
49 'LABEL_44': 'I-CONCEPT',
50 'LABEL_45': 'I-UNKNOWN',
51 'LABEL_46': 'B-EVENT',
52 'LABEL_47': 'B-WEBSITE',
53 'LABEL_48': 'I-NATIONAL',
54 'LABEL_49': 'I-CREATION',
55 'LABEL_50': 'I-ANIMAL'}
56
57model_ckpt = "TTimur/xlm-roberta-base-kyrgyzNER"
58
59config = AutoConfig.from_pretrained(model_ckpt)
60tokenizer = AutoTokenizer.from_pretrained(model_ckpt)
61model = AutoModelForTokenClassification.from_pretrained(model_ckpt, config = config)
62
63# aggregation_strategy = "none"
64nlp = pipeline("ner", model = model, tokenizer = tokenizer, aggregation_strategy = "none")
65
66example = "Кыргызстан Орто Азиянын түндүк-чыгышында орун алган мамлекет."
67ner_results = nlp(example)
68for result in ner_results:
69 result.update({'entity': id2label[result['entity']]})
70 print(result)
71
72# output:
73# {'entity': 'B-LOCATION', 'score': 0.95103735, 'index': 1, 'word': '▁Кыргызстан', 'start': 0, 'end': 10}
74# {'entity': 'B-LOCATION', 'score': 0.79447913, 'index': 2, 'word': '▁Ор', 'start': 11, 'end': 13}
75# {'entity': 'I-LOCATION', 'score': 0.8703734, 'index': 3, 'word': 'то', 'start': 13, 'end': 15}
76# {'entity': 'I-LOCATION', 'score': 0.942387, 'index': 4, 'word': '▁Азия', 'start': 16, 'end': 20}
77# {'entity': 'I-LOCATION', 'score': 0.8542615, 'index': 5, 'word': 'нын', 'start': 20, 'end': 23}
78# {'entity': 'I-LOCATION', 'score': 0.70930535, 'index': 6, 'word': '▁түн', 'start': 24, 'end': 27}
79# {'entity': 'I-LOCATION', 'score': 0.6540094, 'index': 7, 'word': 'дүк', 'start': 27, 'end': 30}
80# {'entity': 'I-LOCATION', 'score': 0.63446337, 'index': 8, 'word': '-', 'start': 30, 'end': 31}
81# {'entity': 'I-LOCATION', 'score': 0.6204858, 'index': 9, 'word': 'чы', 'start': 31, 'end': 33}
82# {'entity': 'I-LOCATION', 'score': 0.6786872, 'index': 10, 'word': 'г', 'start': 33, 'end': 34}
83# {'entity': 'I-LOCATION', 'score': 0.64190257, 'index': 11, 'word': 'ыш', 'start': 34, 'end': 36}
84# {'entity': 'O', 'score': 0.64438057, 'index': 12, 'word': 'ында', 'start': 36, 'end': 40}
85# {'entity': 'O', 'score': 0.9916931, 'index': 13, 'word': '▁орун', 'start': 41, 'end': 45}
86# {'entity': 'O', 'score': 0.9953047, 'index': 14, 'word': '▁алган', 'start': 46, 'end': 51}
87# {'entity': 'O', 'score': 0.9901377, 'index': 15, 'word': '▁мамлекет', 'start': 52, 'end': 60}
88# {'entity': 'O', 'score': 0.99605453, 'index': 16, 'word': '.', 'start': 60, 'end': 61}
89
90
91token = ""
92label_list = []
93token_list = []
94
95for result in ner_results:
96 if result["word"].startswith("▁"):
97 if token:
98 token_list.append(token.replace("▁", ""))
99 token = result["word"]
100 label_list.append(result["entity"])
101 else:
102 token += result["word"]
103
104token_list.append(token.replace("▁", ""))
105
106for token, label in zip(token_list, label_list):
107 print(f"{token}\t{label}")
108
109
110# output:
111# Кыргызстан B-LOCATION
112# Орто B-LOCATION
113# Азиянын I-LOCATION
114# түндүк-чыгышында I-LOCATION
115# орун O
116# алган O
117# мамлекет. O
118
119# aggregation_strategy = "simple"
120nlp = pipeline("ner", model = model, tokenizer = tokenizer, aggregation_strategy = "simple")
121example = "Кыргызстан Орто Азиянын түндүк-чыгышында орун алган мамлекет."
122
123ner_results = nlp(example)
124for result in ner_results:
125 result.update({'entity_group': id2label[result['entity_group']]})
126 print(result)
127
128# output:
129# {'entity_group': 'B-LOCATION', 'score': 0.87275827, 'word': 'Кыргызстан Ор', 'start': 0, 'end': 13}
130# {'entity_group': 'I-LOCATION', 'score': 0.73398614, 'word': 'то Азиянын түндүк-чыгыш', 'start': 13, 'end': 36}
131# {'entity_group': 'O', 'score': 0.92351407, 'word': 'ында орун алган мамлекет.', 'start': 36, 'end': 61}
132