import gradio

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("shkna1368/HemenGI")
model = AutoModelForSeq2SeqLM.from_pretrained("shkna1368/HemenGI")

def remove_html_tags(text):
    """Remove html tags from a string"""
    import re
    clean = re.compile('<.*?>')
    return re.sub(clean, '', text)

def generate(prompt):
     input_ids = tokenizer.encode(prompt, return_tensors="pt")
     #output_ids = model.generate(input_ids, max_length=800, num_beams=20, early_stopping=True)
     complete_str=""
     for i in range(0, 2):
       outputs = model.generate(input_ids, do_sample=True, max_new_tokens=2500,early_stopping=False)
       inp = tokenizer.decode(outputs[0])
       complete_str=complete_str+inp+'\n'+"------"+'\n'

     clear_text=remove_html_tags(complete_str)
     return clear_text

demo=gradio.Interface (
                fn=generate,
                inputs=gradio.Textbox(value="دەتوانی بە شێوازی هێمن شیعری ڕێبەندان و شاعیر لە کتێبی تاریک و ڕوون شیعرێکم بۆ بنووسیت",lines=2, placeholder="شیعرێک", rtl=True),
                outputs=gradio.Textbox(lines=2, placeholder="",rtl=True),
                title='Kurdish Poem Generator',
                description='Kurdish poem generator '
                 
                 )





demo.launch(share=True)