import gradio as gr
import torch
from transformers import pipeline

generate_text = pipeline(
    model="databricks/dolly-v2-12b",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
)


def text_generation(prompt: str):
    return generate_text(prompt)


iface = gr.Interface(fn=text_generation, inputs="text", outputs="text")
iface.launch()