import hf_olmo

from transformers import pipeline
import gradio as gr


olmo_pipe = pipeline("text-generation", model="allenai/OLMo-1B")


def generate(input):
    output = olmo_pipe(input)
    return output[0]["generated_text"]


demo = gr.Interface(fn=generate, 
                    inputs=gr.Textbox(label="Prompt"), 
                    outputs=gr.Textbox(label="Completion"),
                    title="Text Generation with OLMo-1B",
                    description="Generate any text using the `allenai/OLMo-1B` model under the hood!",
                    examples=["Large Language Model is", "The meaning of life is"]
                    )
demo.launch()