Views
No views yet
1from transformers import RobertaForTokenClassification, AutoTokenizer
2model = RobertaForTokenClassification.from_pretrained("guymorlan/levanti_arabic2diacritics")
3tokenizer = AutoTokenizer.from_pretrained("guymorlan/levanti_arabic2diacritics")
4
5label2diacritic = {0: 'ّ', # SHADDA
6 1: 'َ', # FATHA
7 2: 'ِ', # KASRA
8 3: 'ُ', # DAMMA
9 4: 'ْ'} # SUKKUN
10
11
12def arabic2diacritics(text, model, tokenizer):
13 tokens = tokenizer(text, return_tensors="pt")
14 preds = (model(**tokens).logits.sigmoid() > 0.5)[0][1:-1] # remove preds for BOS and EOS
15 new_text = []
16 for p, c in zip(preds, text):
17 new_text.append(c)
18 for i in range(1, 5):
19 if p[i]:
20 new_text.append(label2diacritic[i])
21 # check shadda last
22 if p[0]:
23 new_text.append(label2diacritic[0])
24
25 new_text = "".join(new_text)
26 return new_text
27
28text = "بديش اروح عالمدرسة بكرا"
29arabic2diacritics(text, model, tokenizer)Out[1]: 'بِدِّيْش اْرُوْح عَالْمَدْرَسِة بُكْرَا'