import gradio as gr
import random
from transformers import pipeline

fill_mask = pipeline(
    "fill-mask",
    model="alger-ia/dziribert",
    tokenizer="alger-ia/dziribert"
)

def greet(prompt="أنا", maxTokens=50):
    generated_text_set = set()
    text = prompt +" [MASK]"
    for i in range(int(maxTokens)):
        result = fill_mask(text)
        print(text)
        print(result)
        print(result[0])
        print(result[0][0]['sequence'])
        if result:
            r = random.choice(result[0])['sequence']
            if r not in generated_text_set:
                print(r)
                text = r+ " [MASK]"
                generated_text_set.add(r)
            else:
                continue
    return text

iface = gr.Interface(fn=greet, inputs=["text", "number"], outputs="text", examples=[["أنا [MASK]",50]] )
iface.launch()