import gradio as gr
from inductor import BartInductor

inductor = BartInductor()

def bart(text):
    results = inductor.generate(text, return_scores=True)
    results = [(result[0], float(result[1]) * 100) for result in results]
    results_dict = {result[0]: float(result[1]) for result in results}
    return results_dict

bivar = gr.Interface(fn=bart, 
                     inputs=gr.inputs.Textbox(default='<mask> is the capital of <mask>.'),
                     outputs=gr.Label(),
                     title="BART Inductor 双变量",
                     examples=[['<mask> is the capital of <mask>.'],['<mask> is founder and CEO of <mask>.']],
                     description="Enter a text prompt to generate text using BART.")

trivar = gr.Interface(fn=bart, 
                     inputs=gr.inputs.Textbox(default='<mask> is the parent of <mask>, who is the sibling of <mask>.'),
                     outputs=gr.Label(),
                     title="BART Inductor 三变量",
                     examples=[['<mask> is the parent of <mask>, who is the sibling of <mask>.']],
                     description="this might not work so well")

demo = gr.TabbedInterface([bivar, trivar], ["双变量", "三变量"])

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