from nubia_score import Nubia
import gradio

nubia = Nubia()


def predict(inp_1, inp_2):
    features = nubia.score(inp_1, inp_2, get_features=True)
    labels = {k: v for k, v in features["features"].items()}
    return {"nubia_score": features["nubia_score"]}, labels

title = "NUBIA"
description = """
## A **N**e**u**ral **B**ased **I**nterchangeability **A**ssessor. This is a demo of NUBIA: a SoTA evaluation metric for text generation. Simply input your texts (or click one of the examples to load them) and the model will compute score for how interchangeable they are. Check out the [paper](https://arxiv.org/abs/2004.14667), [blog post](https://wl-research.github.io/blog/), [FAQ](https://github.com/wl-research/nubia/blob/master/FAQ.md) and [demo colab notebook](https://colab.research.google.com/drive/1_K8pOB8fRRnkBPwlcmvUNHgCr4ur8rFg).
"""
inputs = [gradio.inputs.Textbox(label="First Text"), gradio.inputs.Textbox(label="Second Text")]
outputs = [gradio.outputs.Label(label="Interchangeability Score"), gradio.outputs.JSON(label="All Features")]
css =""".gradio-bg[theme=huggingface] .gradio-interface .output-label .confidence {
    color: white !important;
}
.description a {
    text-decoration: underline !important;
}
"""
examples = [
    ["This car is expensive! I can't buy it.", "That automobile costs a fortune! Purchasing it? Impossible!"],
    ["This car is expensive! I can't buy it.", "That automobile costs a good amount. Purchasing it? Totally feasible!"],
    ["The dinner was delicious.", "The dinner did not taste good."]
]
iface = gradio.Interface(fn=predict, inputs=inputs, outputs=outputs, capture_session=True, examples=examples,
                         title=title, description=description, allow_flagging=False, css=css)

iface.launch()