Ahsen Khaliq commited on
Commit
0d34d92
·
1 Parent(s): 9297f2e

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +29 -0
app.py ADDED
@@ -0,0 +1,29 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoTokenizer, AutoModelForCausalLM
4
+
5
+ tokenizer = AutoTokenizer.from_pretrained(
6
+ 'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b',
7
+ bos_token='[BOS]', eos_token='[EOS]', unk_token='[UNK]', pad_token='[PAD]', mask_token='[MASK]'
8
+ )
9
+
10
+ model = AutoModelForCausalLM.from_pretrained(
11
+ 'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b',
12
+ pad_token_id=tokenizer.eos_token_id,
13
+ torch_dtype=torch.float16, low_cpu_mem_usage=True
14
+ ).to(device='cpu', non_blocking=True)
15
+ _ = model.eval()
16
+
17
+
18
+
19
+ def greet(text):
20
+ prompt = text
21
+ with torch.no_grad():
22
+ tokens = tokenizer.encode(prompt, return_tensors='pt').to(device='cpu', non_blocking=True)
23
+ gen_tokens = model.generate(tokens, do_sample=True, temperature=0.8, max_length=64)
24
+ generated = tokenizer.batch_decode(gen_tokens)[0]
25
+
26
+ return generated
27
+
28
+ iface = gr.Interface(fn=greet, inputs="text", outputs="text")
29
+ iface.launch()