import gradio as gr
from transformers import pipeline


model_name = "bigscience/bloom-560m"
nlp_model = pipeline("text-generation", model=model_name, tokenizer=model_name)

def generate_text(input_text, max_length):
    generated_text = nlp_model(input_text, max_length=max_length)
    return f"<br>---<br>".join([sequence["generated_text"] for sequence in generated_text])


input_textbox = gr.inputs.Textbox(lines=5, placeholder="Enter your text here...")
max_length_slider = gr.inputs.Slider(minimum=10, maximum=200, default=50, step=1, label="Max Length")

output_html = gr.outputs.HTML()


gr.Interface(
    fn=generate_text,
    inputs=[input_textbox, max_length_slider],
    outputs=output_html,
    title="Bloom-560m Text Generation",
    description="A demo for the bigscience/bloom-560m model.",
    disable_cache=True,
).launch()