Update README.md

README.md

@@ -30,16 +30,17 @@ label2diacritic = {0: 'ّ', 1: 'َ', 2: 'ِ', 3: 'ُ', 4: ''}
 
 def arabic2diacritics(text, model, tokenizer):
     tokens = tokenizer(text, return_tensors="pt")
-    preds = (model(**tokens).logits.sigmoid() > 0.5)[0]
+    preds = (model(**tokens).logits.sigmoid() > 0.5)[0][1:-1] # remove CLS and SEP
     new_text = []
     for p, c in zip(preds, text):
+        new_text.append(c)
         for i in range(1, 5):
             if p[i]:
                 new_text.append(label2diacritic[i])
         # check shadda last
         if p[0]:
             new_text.append(label2diacritic[0])
-        new_text.append(c)
+        
     new_text = "".join(new_text)
     return new_text