import gradio as gr
from transformers import pipeline

import gc

# Download models
u = pipeline('fill-mask', model='Daniel-Saeedi/Sent-Debias-bert-gender-debiased')
del u
gc.collect()

def fill_mask(stmt,model):
    if model == 'bert':
        debiased = 'Daniel-Saeedi/debiased-bert-base-uncased'
        original = 'bert-base-uncased'

    
    unmasker_1 = pipeline('fill-mask', model=debiased)
    
    unmasker_2 = pipeline('fill-mask', model=original)
    
    return unmasker_1(stmt), unmasker_2(stmt)


demo = gr.Interface(
        fill_mask,
        inputs = [
        gr.Textbox(placeholder="Fill Mask"),
        gr.Radio(choices=['bert'],value='bert')
        ],
        outputs = [gr.Textbox(label="Debiased:"),gr.Textbox(label="Original"), gr.Markdown(
            value="<h3>How the difference is measured?</h3> <p>abs(similarity(gendered_word1,occupation)-similarity(gendered_word2,occupation))</p>")],
            description = '<a href="https://aclanthology.org/2020.acl-main.484/">Double-Hard Debias: Tailoring Word Embeddings for Gender Bias Mitigation</a>'
    )
if __name__ == '__main__':
    demo.launch()