import gradio as gr
from transformers import pipeline

unmasker = pipeline("fill-mask", model="anferico/bert-for-patents")


def unmask(text):
    res = unmasker(text)
    out = {item["token_str"]: item["score"] for item in res}
    return out


textbox = gr.Textbox(label="Type patent abstract here", lines=5)

demo = gr.Interface(
    fn=unmask,
    inputs=textbox,
    outputs="label",
    examples=[
        [
            "The present [MASK] provides a torque sensor that is small and highly rigid and for which high production efficiency is possible."
        ],
        [
            "The present invention relates to [MASK] accessories and pertains particularly to a brake light unit for bicycles."
        ],
        [
            "The present invention discloses a space-bound-free [MASK] and its coordinate determining circuit for determining a coordinate of a stylus pen."
        ],
        [
            "The illuminated [MASK] includes a substantially translucent canopy supported by a plurality of ribs pivotally swingable towards and away from a shaft."
        ],
    ],
)

demo.launch()