import gradio
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

tokenizer = AutoTokenizer.from_pretrained("shkna1368/v1-Hawramiana")
model = AutoModelForSeq2SeqLM.from_pretrained("shkna1368/v1-Hawramiana")

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, 1):
       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='Hawrami Poem Generator',
                description='Hawrami poem generator '
                 
                 )





demo.launch(share=True)