import gradio as gr
from transformers import pipeline

import gc

# Download models
bert_debiased = pipeline('fill-mask', model='Daniel-Saeedi/Sent-Debias-bert-gender-debiased')
bert_original = pipeline('fill-mask', model='bert-base-uncased')

def make_slider(unmask):
    html = '<div><ol>'

    for word in unmask:
        html += '<li><b>{}</b> - Score: {}<li>'.format(word['token_str'],word['score'])
    
    html += '</ol></div>'
    return html
    


def fill_mask(stmt,model):
    if model == 'bert':
        return "<h2>Debiased:</h2>" + make_slider(bert_debiased(stmt)) + "<h2>Original:</h2>" + make_slider(bert_original(stmt))


demo = gr.Interface(
        fill_mask,
        inputs = [
        gr.Textbox(placeholder="Fill Mask"),
        gr.Radio(choices=['bert'],value='bert')
        ],
        outputs = [gr.Markdown(
            value="<h3>Example: </h3> <p>The woman works as [MASK].</p>")],
            description = '<a href="https://arxiv.org/abs/2007.08100">Towards Debiasing Sentence Representations</a>'
    )
if __name__ == '__main__':
    demo.launch()