Create diacritics.py

diacritics.py

@@ -0,0 +1,39 @@
+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