lyric-buddy / app.py
nschenone's picture
Added rap distil model
751b8ae
raw
history blame
849 Bytes
import gradio as gr
from transformers import pipeline, set_seed
chef = pipeline("text-generation", "nschenone/rap-distil")
def generate(text):
max_length: int = 100
num_beams: int = 5
num_return_sequences: int = 1
no_repeat_ngram_size: int = 3
early_stopping: bool = True
skip_special_tokens: bool = True
temperature: float = 1.5
set_seed(0)
generated = chef(
text_inputs=text,
max_length=max_length,
num_beams=num_beams,
num_return_sequences=num_return_sequences,
no_repeat_ngram_size=no_repeat_ngram_size,
early_stopping=early_stopping,
skip_special_tokens=skip_special_tokens,
temperature=temperature
)
return [i["generated_text"] for i in generated]
iface = gr.Interface(fn=generate, inputs="text", outputs="text")
iface.launch()