File size: 374 Bytes
ebce5cd
fb1c9eb
 
ebce5cd
fb1c9eb
 
 
 
 
 
ebce5cd
fb1c9eb
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
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()