import gradio as gr
import nltk
nltk.download('omw-1.4')
from qanom.nominalization_detector import NominalizationDetector

detector = NominalizationDetector()

title = "Nominalization Detection Demo"
description = f"""This is a demo of QANom's nominalization detection algorithm, 
            comprised of candidate nominalization extraction followed by a contextualized binary classification model."""
links = """<p style='text-align: center'>
<a href='https://github.com/kleinay/QANom' target='_blank'>QANom repo</a>  |  
<a href='https://huggingface.co/kleinay/nominalization-candidate-classifier' target='_blank'>Model Repo at Huggingface Hub</a>  |  
<a href='https://www.aclweb.org/anthology/2020.coling-main.274/' target='_blank'>QANom Paper</a>   
</p>"""
examples = [["The doctor was interested in Luke 's treatment .", True, 0.5], 
            ["The description of the horse 's jump provided a surprise to the owner and a show of the skill of the trainer .", True, 0.5],
            ["The construction of the officer 's building finished right after the beginning of the destruction of the previous construction .", True, 0.75]]

def call(sentence: str, return_all_candidates: bool, threshold: float):
    ret = detector([sentence], return_all_candidates, True, threshold)[0]
    if return_all_candidates:
        positives = [d["predicate_idx"] for d in ret if d['predicate_detector_prediction']]
        negatives = [d["predicate_idx"] for d in ret if not d['predicate_detector_prediction']]
    else:
        positives = [d["predicate_idx"] for d in ret]
        negatives = []
    def color(idx):
        if idx in positives: return "lightgreen"
        if idx in negatives: return "pink"
    idx2verb = {d["predicate_idx"] : d["verb_form"] for d in ret}
    idx2prob = {d["predicate_idx"] : d["predicate_detector_probability"] for d in ret}
    def word_span(word, idx):
        tooltip = f'title=" probability={idx2prob[idx]:.2}&#010;verb={idx2verb[idx]}"' if idx in idx2verb else ''
        return f'<span {tooltip} style="background-color: {color(idx)}">{word}</span>'
    html = '<span>' + ' '.join(word_span(word, idx) for idx, word in enumerate(sentence.split(" "))) + '</span>'
    return html, ret

iface = gr.Interface(call, 
                     inputs=[gr.inputs.Textbox(label="Sentence", lines=3), 
                             gr.inputs.Checkbox(default=True, label="Return all candidates?"),
                             gr.inputs.Slider(minimum=0., maximum=1., step=0.01, default=0.5, label="Threshold")], 
                     outputs=[gr.outputs.HTML(label="Detected Nominalizations"), 
                              gr.outputs.JSON(label="Raw Model Output")],
                     title=title,
                     description=description,
                     article=links,
                     examples=examples )
iface.launch()