Update app.py

app.py

@@ -31,13 +31,12 @@ example_1 = [
 ]
 
 example_2 = [
-
-    """به گزارش شانا، تهیه این سند از نیمه دوم پارسال در دستور کار معاونت ضوابط فنی و مهندسی اداره کل نظام فنی و اجرایی و ارزشیابی طرح‌ها قرار گرفت و با برگزاری جلسات تخصصی کارشناسی و استفاده از نظرات حوزه‌های کارفرمایی و پیمانکاری در زمینه پروژه‌های تعمیرات اساسی ماشین‌آلات فرآیندی دوار، کلیات خدمات این حوزه تدوین شد.
- شرکت‌ها و مدیریت‌های زیر مجموعه وزارت نفت که در حوزه‌های تعمیرات اساسی این نوع ماشین‌آلات (شامل انواع توربین، کمپرسور، پمپ، الکتروموتور، ژنراتور و ….) فعالیت می‌کنند، بر اساس این سند قادر خواهند بود شرح خدمات قراردادهای تعمیراتی خود را تهیه کنند.
- این سند پس از ۶ ماه با دریافت دیدگاه‎های پیشنهادی، تکمیل و اصلاح شده و به صورت قطعی ابلاغ می‎شود.
-"""
+"""به گزارش خبرنگار مهر، سرهنگ عین الله جهانی با اشاره به آخرین وضعیت محورهای شمالی اظهار کرد: تردد در محورهای چالوس، فیروزکوه، آزادراه تهران - شمال و آزادراه قزوین-رشت (رفت و برگشت) روان گزارش شده است. 
+ به گفته معاون اجتماعی پلیس راهور ناجا محورهای شمالی دارای بارش باران هستند. 
+ وی با اشاره به آخرین وضعیت ترافیکی سایر محورها، افزود: ترافیک در آزادراه قزوین - کرج - تهران حدفاصل ساسانی تا پل کلاک نیمه سنگین گزارش شده است. 
+ سرهنگ جهانی با اشاره به آخرین وضعیت جوی محورهای مواصلاتی گفت: بارش باران در برخی از محورهای استان‌های آذربایجان شرقی، اردبیل، البرز، خراسان شمالی، گیلان، زنجان، قزوین، تهران، مازندران، سمنان، گلستان و خراسان شمالی گزارش شده است."""
     ,
-    """نمایندگان مجلس شورای اسلامی تصمیم گرفتند که وزارتخانه‌های نفت و صنعت، معدن و تجارت را به انتشار اوراق مالی اسلامی تا سقف ۳۵ هزار میلیارد ریال مجاز کنند، همچنین تصویب شورای اقتصاد برای سرمایه‌گذاری در طرح‌های نفت و گاز اعلام شد."""
+    """معاون اجتماعی پلیس راهور ناجا گفت: تردد در محورهای چالوس، فیروزکوه، آزادراه تهران-شمال، روان گزارش شده است."""
 ]
 
 # Hazm normalizer
@@ -56,14 +55,12 @@ def get_hallucination_label(document, 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"
+    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="text", examples=[example_1, example_2],
+                    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."
                     )