import gradio as gr
from transformers import pipeline
from transformers import pipeline
import re
from utils import grammatical_cleaning

dante = pipeline('text-generation',model='.', tokenizer='GroNLP/gpt2-small-italian-embeddings')

title = "Dante's style text generation"
description = """
Play with Dante's style text generation. Base model is gpt-2-small-italian, fine-tuned on Dante Alighieri's corpus. 
"""
examples = ["Io mi volsi ", "Quand'anche ch'io volli "]

def get_text(input):
	generated = dante(input, max_length=128)[0]['generated_text']
	generated = grammatical_cleaning(generated)
	return generated

iface = gr.Interface(fn=get_text, inputs="text", outputs="text",title=title, examples=examples, description=description)
iface.launch()