import spaces
from transformers import pipeline
import gradio as gr

pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, device='cuda:0')

INSTRUCTION_KEY = "### Instruction:"
RESPONSE_KEY = "### Response:"
INTRO_BLURB = "Below is an instruction that describes a task. Write a response that appropriately completes the request."
PROMPT_FOR_GENERATION_FORMAT = """{intro}
{instruction_key}
{instruction}
{response_key}
""".format(
    intro=INTRO_BLURB,
    instruction_key=INSTRUCTION_KEY,
    instruction="{instruction}",
    response_key=RESPONSE_KEY,
)

example = "James decides to run 3 sprints 3 times a week. He runs 60 meters each sprint. How many total meters does he run a week? Explain before answering."
fmt_ex = PROMPT_FOR_GENERATION_FORMAT.format(instruction=example)

@spaces.GPU
def run():
    with torch.autocast('cuda', dtype=torch.bfloat16):
        return(
            pipe('Here is a recipe for vegan banana bread:\n',
                max_new_tokens=100,
                do_sample=True,
                use_cache=True))



with gr.Blocks() as app:
    btn = gr.Button()
    outp=gr.Textbox()
    btn.click(run,None,outp)
app.launch()