File size: 769 Bytes
a88d1fa
 
c72147b
 
 
 
 
 
 
 
 
 
a88d1fa
 
 
 
05bb1dc
e09a7c9
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from instruct_pipeline import InstructionTextGenerationPipeline

tokenizer       = AutoTokenizer.from_pretrained("databricks/dolly-v2-3b", 
                                          padding_side="left")
model           = AutoModelForCausalLM.from_pretrained("databricks/dolly-v2-3b", 
                                             torch_dtype=torch.bfloat16)
                                             
generate_text   = InstructionTextGenerationPipeline(model=model, 
                                                    tokenizer=tokenizer)

def textGen(text):
    return generate_text(text)
    
demo = gr.Interface(fn=textGen, inputs="text", outputs="text")
demo.launch()