import gradio as gr
from qasrl_model_pipeline import QASRL_Pipeline

models = ["kleinay/qanom-seq2seq-model-baseline", 
          "kleinay/qanom-seq2seq-model-joint",
          "kleinay/qasrl-seq2seq-model"]
pipelines = {model: QASRL_Pipeline(model) for model in models}


description = f"""This is a demo of QASRL/QANom models, which fine-tuned a Seq2Seq pretrained model (T5) on the QASRL/QANom tasks.""" 
title="QANom/QASRL Parser Demo"
examples = [[models[0], "The doctor was interested in Luke 's <p> treatment .", True, "treat"],
            [models[1], "The doctor was interested to know about Luke 's bio-feedback <p> treatment given by the nurse yesterday.", True, "treat"],
            [models[2], "The doctor was interested to <p> know about Luke 's bio-feedback treatment given by the nurse yesterday.", False, "know"],
            [models[1], "The Veterinary student was interested in Luke 's <p> treatment of sea animals .", True, "treat"],
            [models[1], "The Veterinary student was <p> interested in Luke 's treatment of sea animals .", False, "interest"]]

input_sent_box_label = "Insert sentence here. Mark the predicate by adding the token '<p>' before it."
verb_form_inp_placeholder = "e.g. 'decide' for the nominalization 'decision', 'teach' for 'teacher', etc."
links = """<p style='text-align: center'>
<a href='https://www.qasrl.org' target='_blank'>QASRL Website</a>  |  <a href='https://huggingface.co/kleinay/qanom-seq2seq-model-baseline' target='_blank'>Model Repo at Huggingface Hub</a>
</p>"""
def call(model_name, sentence, is_nominal, verb_form):
    predicate_marker="<p>"
    if predicate_marker not in sentence:
        raise ValueError("You must highlight one word of the sentence as a predicate using preceding '<p>'.")
        
    if not verb_form:
        if is_nominal:
            raise ValueError("You should provide the verbal form of the nominalization")
            
        toks = sentence.split(" ")
        pred_idx = toks.index(predicate_marker)
        predicate = toks(pred_idx+1)
        verb_form=predicate
    pipeline = pipelines[model_name]
    pipe_out = pipeline([sentence], 
                    predicate_marker=predicate_marker, 
                    predicate_type="nominal" if is_nominal else "verbal",
                    verb_form=verb_form)[0]
    return pipe_out["QAs"], pipe_out["generated_text"]
iface = gr.Interface(fn=call, 
                     inputs=[gr.inputs.Radio(choices=models, default=models[0], label="Model"), 
                             gr.inputs.Textbox(placeholder=input_sent_box_label, label="Sentence", lines=4), 
                             gr.inputs.Checkbox(default=True, label="Is Nominalization?"),
                             gr.inputs.Textbox(placeholder=verb_form_inp_placeholder, label="Verbal form (for nominalizations)", default='')], 
                     outputs=[gr.outputs.JSON(label="Model Output - QASRL"), gr.outputs.Textbox(label="Raw output sequence")],
                     title=title,
                     description=description,
                     article=links,
                     examples=examples )
                     
iface.launch()