import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
import transformers
import hazm
import gradio as gr

# Define class of the model
class ParsbertHallu(nn.Module, PyTorchModelHubMixin):
    def __init__(self):
        super().__init__()
        self.transformer_model = transformers.AutoModelForSequenceClassification.from_pretrained("Pooya-Fallah/ParsBERT-nli-FarsTail-FarSick",
                                                                           num_labels=3)
        self.head = nn.Sequential(
            nn.Linear(3,1),
            nn.Sigmoid()
        )


    def forward(self, x):
        out = self.transformer_model(**x)['logits']
        return torch.squeeze(self.head(out))

# Example Inputs
example_1 = [
    """به گزارش شانا، نمایندگان در نشست علنی عصر دیروز (شنبه، ۱۴ بهمن‌ماه) مجلس شورای اسلامی، در جریان بررسی جزییات لایحه بودجه سال ۹۷ کل کشور با بخش درآمدی بند (ط) تبصره ۵ این لایحه با ۱۵۶ رای موافق، ۲۲ رای مخالف و ۵ رای ممتنع از مجموع ۲۳۱ نماینده حاضر در جلسه موافقت کردند.
 در بخش درآمدی بند (ط) تبصره (۵) این لایحه آمده است که به‌‌منظور‌ سرمایه‌‌گذاری در طرح‌های نفت و گاز با اولویت میدان‌های مشترک وزارت نفت و طرح‌های زیربنایی و توسعه‌ای وزارت صنعت، معدن و تجارت، وزارتخانه‌های مذکور از طریق شرکت‌های تابع ذی‌ربط و با تصویب شورای اقتصاد، اوراق مالی اسلامی (ریالی یا ارزی) در سقف ۳۵ هزار میلیارد ریال منتشر کنند.
"""
,
    """نمایندگان مجلس شورای اسلامی تصمیم گرفتند که وزارتخانه‌های نفت و صنعت، معدن و تجارت را به انتشار اوراق مالی اسلامی تا سقف ۳۵ هزار میلیارد ریال مجاز کنند، همچنین تصویب شورای اقتصاد برای سرمایه‌گذاری در طرح‌های نفت و گاز اعلام شد."""
]

example_2 = [

    """به گزارش شانا، تهیه این سند از نیمه دوم پارسال در دستور کار معاونت ضوابط فنی و مهندسی اداره کل نظام فنی و اجرایی و ارزشیابی طرح‌ها قرار گرفت و با برگزاری جلسات تخصصی کارشناسی و استفاده از نظرات حوزه‌های کارفرمایی و پیمانکاری در زمینه پروژه‌های تعمیرات اساسی ماشین‌آلات فرآیندی دوار، کلیات خدمات این حوزه تدوین شد.
 شرکت‌ها و مدیریت‌های زیر مجموعه وزارت نفت که در حوزه‌های تعمیرات اساسی این نوع ماشین‌آلات (شامل انواع توربین، کمپرسور، پمپ، الکتروموتور، ژنراتور و ….) فعالیت می‌کنند، بر اساس این سند قادر خواهند بود شرح خدمات قراردادهای تعمیراتی خود را تهیه کنند.
 این سند پس از ۶ ماه با دریافت دیدگاه‎های پیشنهادی، تکمیل و اصلاح شده و به صورت قطعی ابلاغ می‎شود.
"""
    ,
    """نمایندگان مجلس شورای اسلامی تصمیم گرفتند که وزارتخانه‌های نفت و صنعت، معدن و تجارت را به انتشار اوراق مالی اسلامی تا سقف ۳۵ هزار میلیارد ریال مجاز کنند، همچنین تصویب شورای اقتصاد برای سرمایه‌گذاری در طرح‌های نفت و گاز اعلام شد."""
]

# Hazm normalizer
normalizer = hazm.Normalizer()

# tokenizer is from ParsBERT (HooshvareLab/bert-fa-zwnj-base)
tokenizer = transformers.AutoTokenizer.from_pretrained('HooshvareLab/bert-fa-zwnj-base')

# load model
model = ParsbertHallu.from_pretrained("Pooya-Fallah/ParsbertHallu")

def get_hallucination_label(document, summary):
    document = normalizer.normalize(document)
    summary = normalizer.normalize(summary)

    tokens = tokenizer(document, summary, truncation=True, padding=True,
                          max_length=512, return_tensors='pt')

    prob = model(tokens).item()
    if prob > 0.5:
        return "Hallucinated"
    else:
        return "Not-Hallucinated"

demo = gr.Interface(fn=get_hallucination_label, inputs=[gr.TextArea(lines=7, placeholder="document"), gr.TextArea(lines=3, placeholder="summary")],
                    outputs="text", examples=[example_1, example_2],
                    title="Hallucination Detection Demo for Persian Summarization Task",
                    description="A straightforward binary classifier that determines whether the generated summarization is hallucinated or not."
                    )
demo.launch()