import gradio as gr
import numpy as np
from transformers import pipeline, Pipeline

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

example = 'A crustless sandwich made from two slices of baked bread'

# class TempScalePipe(Pipeline):
#     def _forward(self, model_inputs):
#         outputs = self.model(**model_inputs)
#         return outputs
#
#     def postprocess(self, model_outputs, temp=1e3):
#         out = model_outputs["logits"] / temp
#         idx = np.random.randint(-3, 0)
#         best_class = out.softmax(idx)
#         return best_class
#

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


textbox = gr.Textbox(label="Type language here", lines=5)
import gradio as gr
from transformers import pipeline, Pipeline


# unmasker = pipeline("fill-mask", model="anferico/bert-for-patents")
#
#
# def add_mask(text, size=1):
#     split_text = text.split()
#     idx = np.random.randint(len(split_text), size=size)
#     for i in idx:
#         split_text[i] = '[MASK]'
#     return ' '.join(split_text)
#
#
# def unmask(text):
#     text = add_mask(text)
#     res = unmasker(text)
#     out = {item["token_str"]: item["score"] for item in res}
#     return out
#
#
# textbox = gr.Textbox(label="Type language here", lines=5)
#
# demo = gr.Interface(
#     fn=unmask,
#     inputs=textbox,
#     outputs="label",
#     examples=[
#
#     ],
# )

demo.launch()
demo = gr.Interface(
    fn=unmask,
    inputs=textbox,
    outputs="label",
    examples=[
    ],
)

demo.launch()