import gradio as gr
from transformers import pipeline
import numpy as np

hhem = pipeline("text-classification", model="vectara/hallucination_evaluation_model")

def get_hhem_score(sentence1, sentence2):
    output = hhem(f'{sentence1} [SEP] {sentence2}')
    score = np.round(output[0]['score'], 4)

    return score

demo = gr.Interface(
    fn=get_hhem_score,
    inputs=[
        gr.components.Textbox(label="Sentence 1"),
        gr.components.Textbox(label="Sentence 2"),
    ],
    outputs=gr.components.Label(num_top_classes=1, label='HHEM Score'),
    examples=[
        ["Vectara provides RAG-as-a-service", "RAG-as-a-service is provided by Vectara"],
        ["The cat sat on the mat.", "A feline was resting on a small rug."],
        ["The quick brown fox jumps over the lazy dog.", "A fast red fox leaps across a sleepy canine."],
    ],
    cache_examples=False,
    allow_flagging="never",
    flagging_options=None,
    title="HHEM Demo",
    description="This demo uses Vectara's Hallucination Evaluation model (HHEM) to calculate factual consistency between two input sentences.",
)

# Launch the demo
demo.launch()