File size: 1,366 Bytes
c3d652a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch
import hazm
# بارگذاری مدل ParsBERT
model_name = "HooshvareLab/bert-fa-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(model_name)
def add_diacritics(text):
# نرمالسازی و توکنسازی
normalizer = hazm.Normalizer()
text = normalizer.normalize(text)
words = hazm.word_tokenize(text)
# پردازش ورودی برای مدل
inputs = tokenizer(words, return_tensors="pt", is_split_into_words=True)
# پیشبینی مدل
with torch.no_grad():
outputs = model(**inputs).logits
# دریافت لیبلهای پیشبینیشده
predictions = torch.argmax(outputs, dim=2).tolist()[0]
# قوانین اضافه کردن اعراب
diacritics = {1: 'َ', 2: 'ِ', 3: 'ُ'} # فتحه، کسره، ضمه
result = []
for word, prediction in zip(words, predictions):
if prediction in diacritics:
word += diacritics[prediction]
result.append(word)
# بازسازی جمله با رعایت علائم نگارشی
final_text = " ".join(result)
final_text = final_text.replace(" ،", "،").replace(" .", ".").replace(" ؛", "؛")
return final_text
|