import gc 
import gradio as gr
from transformers import pipeline

pipe = pipeline('text-generation', model='akhooli/ap2023', tokenizer='akhooli/ap2023')
#gc.collect()

def sayPoetry(prompt):
    gen = pipe(prompt, max_length=96)[0]["generated_text"]
    poetry =""
    for line in gen.split('.')[:-1]:
        poetry += line
    return poetry
                
tb = gr.Textbox(label="Enter your prompt or select from examples:", placeholder="أنا الذي", lines=1)
poetry = gr.Interface(fn=sayPoetry, inputs=tb, outputs="text", allow_flagging='never',
                   title='Arabic Poetry Generation Demo (updated Jan. 2023)',
                   description = "a simple demo of AI generated poetry based on 1M poems fine-tuned using AraGPT2 (be patient, runs on cpu)",
                   examples=[['أنت'],['لولا كتاب'],['ألا ليت'],['يا قدس'],['عيد بأية حال'],['لكل شيء إذا ما'],['.']])

poetry.launch()