import gradio as gr
from transformers import pipeline

generator = pipeline(
    "text-generation",
    model="cyberagent/open-calm-small",
    tokenizer="cyberagent/open-calm-small",
)


def generate(text):
    result = generator(text)
    return result[0]["generated_text"]


examples = [
    ["AIによって私達の暮らしは、"],
]

demo = gr.Interface(
    fn=generate,
    inputs=gr.Textbox(lines=5, label="Input Text"),
    outputs=gr.Textbox(lines=5, label="Generated Text"),
    examples=examples,
    description="# [CyberAgent OpenCALM-Small](https://huggingface.co/cyberagent/open-calm-small)",
)

demo.launch()