app.py

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 = round(model(tokens).item(), 3)

    return {"Hallucinated": prob, "Not-Hallucinated": 1-prob}

demo = gr.Interface(fn=get_hallucination_label, inputs=[gr.TextArea(lines=7, placeholder="document"), gr.TextArea(lines=3, placeholder="summary")],
                    outputs=gr.Label(num_top_classes=2), 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()