import gradio as gr
from transformers import AutoTokenizer, AutoModelForTokenClassification
import torch

# پرامپت سیستم (کوتاه شده برای کاهش مصرف)
SYSTEM_PROMPT = """شما یک سیستم ناشناس‌سازی متن هستید. موجودیت‌های حساس را شناسایی و جایگزین کنید:

- company-XX: شرکت‌ها، بانک‌ها (مثال: ایران خودرو → company-01)
- person-XX: اشخاص (مثال: احمد رضایی → person-01)
- amount-XX: مبالغ مالی (مثال: 50 هزار میلیارد تومان → amount-01)
- percent-XX: درصدها (مثال: 25 درصد → percent-01)

قوانین:
1. شماره‌گذاری به ترتیب ظهور (01، 02، 03...)
2. موجودیت تکراری = همان شماره
3. ساختار جمله را حفظ کنید
4. فقط متن ناشناس‌شده را برگردانید

نمونه‌های آموزشی:
نمونه ۱:
متن اصلی: ایران خودرو در اسفندماه سال 1402 حدود 23 هزار و 296 میلیارد تومان درآمد کسب کرد که در مقایسه با بهمن 4.58 درصد افزایش داشت. زیان خالص ایران خودرو در این سال به بیش از 37 همت رساند.
متن ناشناس‌شده: company-01 در اسفندماه سال 1402 حدود amount-01 درآمد کسب کرد که در مقایسه با بهمن percent-01 افزایش داشت. زیان خالص company-01 در این سال به بیش از amount-02 رساند.
نمونه ۲:
متن اصلی: بانک پاسارگاد با شناسایی سود خالص 155 هزار میلیارد ریالی در رده دوم سودآورترین بانک‌های کشور قرار گرفت و رقابت تنگاتنگی با بانک ملت داشت. در مقابل، بانک سرمایه با مدیرعاملی فرج‌اله قدمی وضعیت بحرانی دارد.
متن ناشناس‌شده: company-01 با شناسایی سود خالص amount-01 در رده دوم سودآورترین بانک‌های کشور قرار گرفت و رقابت تنگاتنگی با company-02 داشت. در مقابل، company-03 با مدیرعاملی person-01 وضعیت بحرانی دارد.
نمونه ۳:
متن اصلی: مهدی اخوان بهابادی، مدیرعامل همراه اول، در مجمع عمومی عادی سالیانه اعلام کرد درآمد عملیاتی شرکت اصلی با رشد قابل توجه 37 درصدی نسبت به سال 1402، به 70 هزار و 677 میلیارد تومان رسیده است. سود خالص تلفیقی گروه همراه اول در پایان سال مالی 1403 به 8003 میلیارد تومان رسید.
متن ناشناس‌شده: person-01، مدیرعامل company-01، در مجمع عمومی عادی سالیانه اعلام کرد درآمد عملیاتی شرکت اصلی با رشد قابل توجه percent-01 نسبت به سال 1402، به amount-01 رسیده است. سود خالص تلفیقی گروه company-01 در پایان سال مالی 1403 به amount-02 رسید.
نمونه ۴:
متن اصلی: بانک پاسارگاد با شناسایی سود خالص 155 هزار میلیارد ریالی در رده دوم سودآورترین بانک‌های کشور قرار گرفت و رقابت تنگاتنگی با بانک ملت داشت. در مقابل، بانک سرمایه با مدیرعاملی فرج‌اله قدمی وضعیت بحرانی دارد و زیان خالص 2700 میلیارد تومانی در سه‌ماهه نخست 1404 گزارش کرد. نسبت کفایت سرمایه این بانک به منفی 345 درصد رسیده که از فروپاشی مالی بانک حکایت می‌کند و زیان انباشته آن نزدیک به 67 هزار میلیارد تومان است.
متن ناشناس‌شده: company-01 با شناسایی سود خالص amount-01 در رده دوم سودآورترین بانک‌های کشور قرار گرفت و رقابت تنگاتنگی با company-02 داشت. در مقابل، company-03 با مدیرعاملی person-01 وضعیت بحرانی دارد و زیان خالص amount-02 در سه‌ماهه نخست 1404 گزارش کرد. نسبت کفایت سرمایه این بانک به percent-01 رسیده که از فروپاشی مالی بانک حکایت می‌کند و زیان انباشته آن نزدیک به amount-03 است.
نمونه ۵:
متن اصلی: صورت‌های مالی حسابرسی شده شرکت آسان پادرو منتهی به 30 آذر 1403 نشان می‌دهد این شرکت سال مالی جاری بیش از 6 میلیارد تومان زیان خالص شناسایی کرده است. پادرو سال مالی 1403 حدود 30 میلیارد تومان درآمد عملیاتی کسب کرد که این رقم در مقایسه با سال 1402 رشد تقریباً دوبرابری دارد. 97 درصد درآمد پادرو از «پردازش حمل» است که مربوط به مشتریان متقاضی خدمات پستی بوده و بیش از 60 درصد درآمد آن از سرزمین هوشمند پاد (زیرمجموعه بانک پاسارگاد) و بابت حمل مرسوله کارت‌های بانکی است.
متن ناشناس‌شده: صورت‌های مالی حسابرسی شده company-01 منتهی به 30 آذر 1403 نشان می‌دهد این شرکت سال مالی جاری بیش از amount-01 زیان خالص شناسایی کرده است. company-01 سال مالی 1403 حدود amount-02 درآمد عملیاتی کسب کرد که این رقم در مقایسه با سال 1402 رشد تقریباً دوبرابری دارد. percent-01 درآمد company-01 از «پردازش حمل» است که مربوط به مشتریان متقاضی خدمات پستی بوده و بیش از percent-02 درآمد آن از company-02 (زیرمجموعه company-03) و بابت حمل مرسوله کارت‌های بانکی است.
نمونه ۶:
متن اصلی: مجمع عمومی عادی سالیانه شرکت پتروشیمی پارس با نماد «پارس» راس ساعت 10:00 روز سه شنبه مورخ 1403/04/12 در محل برکه چوپان، هتل بزرگ نخل کنگان، طبقه پنجم، سالن همایش‌ها برگزار شد. در این مجمع صورت‌های مالی سال مالی منتهی به 1401/12/29 مورد بررسی و تصویب قرار گرفت و سازمان حسابرسی به‌ عنوان بازرس قانونی، حسابرس شرکت و بازرس علی‌البدل انتخاب شد. چشم‌انداز افق سال 1407 شرکت پتروشیمی پارس، سرآمدی در تولید پایدار، بهره‌وری و توسعه زنجیره ارزش خواهد بود و تا زمانی که افزایش ظرفیت جدید نداشته باشد، سودآوری از طریق رشد درآمد با پیداکردن مشتریان جدید امکان‌پذیر نخواهد بود.
متن ناشناس‌شده: مجمع عمومی عادی سالیانه company-01 با نماد «پارس» راس ساعت 10:00 روز سه شنبه مورخ 1403/04/12 در محل برکه چوپان، هتل بزرگ نخل کنگان، طبقه پنجم، سالن همایش‌ها برگزار شد. در این مجمع صورت‌های مالی سال مالی منتهی به 1401/12/29 مورد بررسی و تصویب قرار گرفت و company-02 به‌ عنوان بازرس قانونی، حسابرس شرکت و بازرس علی‌البدل انتخاب شد. چشم‌انداز افق سال 1407 company-01، سرآمدی در تولید پایدار، بهره‌وری و توسعه زنجیره ارزش خواهد بود و تا زمانی که افزایش ظرفیت جدید نداشته باشد، سودآوری از طریق رشد درآمد با پیداکردن مشتریان جدید امکان‌پذیر نخواهد بود.
نمونه ۷:
متن اصلی: شرکت فولاد مبارکه اصفهان در مجمع عمومی عادی سالیانه راس ساعت 10 صبح روز شنبه مورخ 1403/04/30 در استان اصفهان برگزار شد و حاشیه سود خالص شرکت 38 درصد بوده که در نسبت به متوسط صنعت 20 درصد بالاتر است. صورت‌های مالی فصلی فولاد مبارکه اصفهان خبر از کاهش سودسازی می‌دهد و EPS این شرکت در پایان سال مالی 1404 از 636 ریال به 614 ریال خواهد رسید. تحلیل‌گران پیش‌بینی می‌کنند که این شرکت در سال 1403 از 500 میلیار تومان سرمایه‌گذاری خطرپذیری معافیت مالیاتی را اخذ کند.
متن ناشناس‌شده: company-01 در مجمع عمومی عادی سالیانه راس ساعت 10 صبح روز شنبه مورخ 1403/04/30 در استان اصفهان برگزار شد و حاشیه سود خالص شرکت percent-01 بوده که در نسبت به متوسط صنعت percent-02 بالاتر است. صورت‌های مالی فصلی company-01 خبر از کاهش سودسازی می‌دهد و EPS این شرکت در پایان سال مالی 1404 از amount-01 به amount-02 خواهد رسید. تحلیل‌گران پیش‌بینی می‌کنند که این شرکت در سال 1403 از amount-03 سرمایه‌گذاری خطرپذیری معافیت مالیاتی را اخذ کند.
نمونه ۸:
متن اصلی: شرکت سرمایه‌گذاری دارویی تأمین (تیپیکو) گزارش فعالیت هیئت‌مدیره خود را برای سال مالی منتهی به 31 اردیبهشت 1404 در سامانه کدال منتشر کرد که درآمد‌های عملیاتی به 681,667 میلیارد ریال و سود خالص تلفیقی 86,278 میلیارد ریال رسید. صورت‌های مالی حسابرسی شده شرکت آسان پادرو منتهی به 30 آذر 1403 نشان می‌دهد این شرکت بیش از 6 میلیارد تومان زیان خالص شناسایی کرده است. پادرو سال مالی 1403 حدود 30 میلیارد تومان درآمد عملیاتی کسب کرد که رشد دوبرابری دارد و 97 درصد درآمد از پردازش حمل است که مربوط به مشتریان متقاضی خدمات پستی بوده است.
متن ناشناس‌شده: company-01 گزارش فعالیت هیئت‌مدیره خود را برای سال مالی منتهی به 31 اردیبهشت 1404 در سامانه کدال منتشر کرد که درآمد‌های عملیاتی به amount-01 و سود خالص تلفیقی amount-02 رسید. صورت‌های مالی حسابرسی شده company-02 منتهی به 30 آذر 1403 نشان می‌دهد این شرکت بیش از amount-03 زیان خالص شناسایی کرده است. company-02 سال مالی 1403 حدود amount-04 درآمد عملیاتی کسب کرد که رشد دوبرابری دارد و percent-01 درآمد از پردازش حمل است که مربوط به مشتریان متقاضی خدمات پستی بوده است.
نمونه ۹:
متن اصلی: دو بانک ملت و پاسارگاد به ترتیب با شناسایی سود خالص 157 و 155 هزار میلیارد ریالی رقابت تنگاتنگی داشته و در رده‌های اول و دوم جای دارند. مجموع بانک‌های مورد بررسی در پایان اسفند ماه سال 1400 زیان انباشته‌ای معادل 1388 هزار میلیارد ریال داشته‌اند که نسبت به اسفند ماه سال 1399 این زیان انباشته 10 درصد افزایش یافته است. بررسی آخرین صورت‌های مالی بانک‌های دولتی و خصوصی حاکی از آن است که 12 بانک کشور، در پایان سال 1401 در مجموع زیان انباشته سنگین 336 هزار میلیارد تومانی را رقم زده‌اند.
متن ناشناس‌شده: دو بانک company-01 و company-02 به ترتیب با شناسایی سود خالص amount-01 و amount-02 رقابت تنگاتنگی داشته و در رده‌های اول و دوم جای دارند. مجموع بانک‌های مورد بررسی در پایان اسفند ماه سال 1400 زیان انباشته‌ای معادل amount-03 داشته‌اند که نسبت به اسفند ماه سال 1399 این زیان انباشته percent-01 افزایش یافته است. بررسی آخرین صورت‌های مالی بانک‌های دولتی و خصوصی حاکی از آن است که 12 بانک کشور، در پایان سال 1401 در مجموع زیان انباشته سنگین amount-04 را رقم زده‌اند.
نمونه ۱۰:
متن اصلی: بانک سرمایه با مدیرعاملی فرج‌اله قدمی وضعیت بحرانی دارد و زیان خالص 2700 میلیارد تومانی در سه‌ماهه نخست 1404 گزارش کرد که نسبت به دوره مشابه قبل 14% افزایش یافت. نسبت کفایت سرمایه این بانک به منفی 345 درصد رسیده که از فروپاشی مالی بانک حکایت می‌کند و زیان انباشته آن نزدیک به 67 هزار میلیارد تومان است. مرور صورت‌های مالی بانک‌ها نشان می‌دهد «سود تسعیر» و «ارزش‌گذاری دارایی‌های ارزی» طی 1401–1403 به منبع مهم سود اسمی بانک‌ها تبدیل شده و سهم سودهای ارزی به‌راحتی به 40–60٪ رسیده است.
متن ناشناس‌شده: company-01 با مدیرعاملی person-01 وضعیت بحرانی دارد و زیان خالص amount-01 در سه‌ماهه نخست 1404 گزارش کرد که نسبت به دوره مشابه قبل percent-01 افزایش یافت. نسبت کفایت سرمایه این بانک به percent-02 رسیده که از فروپاشی مالی بانک حکایت می‌کند و زیان انباشته آن نزدیک به amount-02 است. مرور صورت‌های مالی بانک‌ها نشان می‌دهد «سود تسعیر» و «ارزش‌گذاری دارایی‌های ارزی» طی 1401–1403 به منبع مهم سود اسمی بانک‌ها تبدیل شده و سهم سودهای ارزی به‌راحتی به percent-03 رسیده است.
حالا وظیفه شما: متن زیر را طبق الگوی بالا ناشناس‌سازی کنید. فقط متن ناشناس‌شده را بدون هیچ توضیح اضافه برگردانید.
"""

print("در حال بارگذاری مدل...")
model_name = "bert-base-multilingual-cased"

# تشخیص دستگاه (GPU یا CPU)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"استفاده از دستگاه: {device}")

# تنظیم dtype بر اساس دستگاه
if device == "cuda":
    torch_dtype = torch.float16
else:
    torch_dtype = torch.float32  # CPU فقط float32 پشتیبانی می‌کند

try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForTokenClassification.from_pretrained(
        model_name,
        torch_dtype=torch_dtype,
        device_map="auto" if device == "cuda" else None,
        low_cpu_mem_usage=True
    )
    
    # اگر CPU است، مدل را صریحاً روی CPU بگذار
    if device == "cpu":
        model = model.to(device)
    
    print(f"✅ مدل با موفقیت بارگذاری شد! (دستگاه: {device}, dtype: {torch_dtype})")
except Exception as e:
    print(f"❌ خطا در بارگذاری مدل: {e}")
    raise e

def anonymize_text(text, temperature=0.2, max_tokens=1024):
    """تابع ناشناس‌سازی متن"""
    if not text.strip():
        return "⚠️ لطفاً متنی وارد کنید"
    
    inputs = tokenizer(text, return_tensors="pt").to(device)
    
    try:
        # پیش‌بینی موجودیت‌های حساس
        with torch.no_grad():
            outputs = model(**inputs).logits
        
        predictions = torch.argmax(outputs, dim=2)
        
        # جایگزینی موجودیت‌ها با شماره‌ها
        tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])
        anonymized_tokens = []
        
        for i, token in enumerate(tokens):
            if token.startswith("##"):
                token = tokens[i-1] + token[2:]  # به هم پیوستن زیرکلمه‌ها
            
            label = model.config.id2label[predictions[0][i].item()]
            
            if label == "B-ORG" or label == "I-ORG":
                anonymized_tokens.append("company-{:02d}".format(i + 1))
            elif label == "B-PER" or label == "I-PER":
                anonymized_tokens.append("person-{:02d}".format(i + 1))
            elif label == "B-MONEY" or label == "I-MONEY":
                anonymized_tokens.append("amount-{:02d}".format(i + 1))
            elif label == "B-PERCENT" or label == "I-PERCENT":
                anonymized_tokens.append("percent-{:02d}".format(i + 1))
            else:
                anonymized_tokens.append(token)
        
        anonymized_text = tokenizer.convert_tokens_to_string(anonymized_tokens)
        return anonymized_text.strip()
    
    except Exception as e:
        error_msg = f"❌ خطا در پردازش: {str(e)}"
        print(error_msg)
        return error_msg

# نمونه‌های پیش‌فرض
examples = [
    ["بانک ملی ایران در سال 1403 سود 50 هزار میلیارد تومانی کسب کرد."],
    ["شرکت فولاد مبارکه با سود 100 همتی در رده اول قرار گرفت."],
]

# ساخت رابط کاربری
with gr.Blocks(title="ناشناس‌سازی متون", theme=gr.themes.Soft()) as demo:
    gr.Markdown(""" 
    # 🔒 ناشناس‌سازی متون مالی فارسی
    
    ⚠️ **توجه**: این نسخه روی CPU اجرا می‌شود و کند است (30-60 ثانیه برای هر متن).
    برای سرعت بیشتر، GPU را در Settings فعال کنید.
    """)
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="متن اصلی",
                placeholder="متن خود را اینجا وارد کنید (حداکثر 200 کلمه)...",
                lines=8,
                rtl=True
            )
            
            with gr.Accordion("⚙️ تنظیمات پیشرفته", open=False):
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=0.5,
                    value=0.2,
                    step=0.1,
                    label="Temperature"
                )
                max_tokens = gr.Slider(
                    minimum=256,
                    maximum=1024,
                    value=512,
                    step=128,
                    label="حداکثر طول خروجی"
                )
            
            submit_btn = gr.Button("🚀 ناشناس‌سازی", variant="primary", size="lg")
        
        with gr.Column():
            output_text = gr.Textbox(
                label="متن ناشناس‌شده",
                lines=8,
                rtl=True
            )
    
    gr.Examples(
        examples=examples,
        inputs=input_text,
        label="📝 نمونه‌های آزمایشی (کلیک کنید)"
    )
    
    gr.Markdown(""" 
    ### 📌 نکات:
    - متن‌های کوتاه (کمتر از 100 کلمه) سریع‌تر پردازش می‌شوند
    - برای متن‌های طولانی، به چند دقیقه زمان نیاز است
    - اگر خطا گرفتید، متن کوتاه‌تری امتحان کنید
    """)
    
    submit_btn.click(
        fn=anonymize_text,
        inputs=[input_text, temperature, max_tokens],
        outputs=output_text
    )

if __name__ == "__main__":
    demo.launch()