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