import gradio as gr
from inductor import BartInductor

inductor = BartInductor()

def bart(prompt, num, return_score):
    results = inductor.generate(prompt, k=num, topk=num, return_scores=return_score)
    if return_score:
        results = [(result[0], float(result[1]) * 100) for result in results]
    else:
        results = [(result, 0) for result in results]
    results_dict = {result[0]: float(result[1]) for result in results}
    return results_dict

demo = gr.Interface(fn=bart, 
                     inputs=[gr.inputs.Textbox(default='<mask> is the capital of <mask>.'),
                             gr.Slider(0, 10, value=5, step=1),
                             gr.Checkbox(label="Hell Yes", info="Return Scores?")],
                     outputs=gr.Label(),
                     title="ORION: Open Rule InductiON",
                     examples=[['<mask> is the capital of <mask>.', 5, True],
                               ['<mask> is founder and CEO of <mask>.', 5, False],
                              ["<mask>'s mother was a <mask>-based actress, <mask>.", 5, False]],
                     description="Enter a text prompt to generate text. Make sure using <mask> to replace entities.")


if __name__ == "__main__":
    demo.launch()