jason9693's picture
modified some arguments
621210f
raw
history blame
2 kB
import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, StoppingCriteriaList
import os
import torch
theme = "darkgrass"
title = "GPT-NeoX(Korean) Demo"
model_name = "EleutherAI/gpt-neox-ko-1.3b"
description = "GPT-NeoX ν•œκ΅­μ–΄ λͺ¨λΈμ„ μ‹œμ—°ν•˜λŠ” 데λͺ¨νŽ˜μ΄μ§€ μž…λ‹ˆλ‹€."
# article = "<p style='text-align: center'><a href='https://github.com/kingoflolz/mesh-transformer-jax' target='_blank'>GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model</a></p>"
examples = [
["μΈκ°„μ²˜λŸΌ μƒκ°ν•˜κ³ , ν–‰λ™ν•˜λŠ” '지λŠ₯'을 톡해"],
["질문: ν•œκ°€μœ„λ₯Ό κ°€μž₯ 의미있게 λ³΄λ‚Όμˆ˜ μžˆλŠ” 방법은 λ¬΄μ—‡μΌκΉŒμš”? \nλ‹΅λ³€:"],
["2040λ…„ 미ꡭ은, "]
]
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-ko-1.3b", use_auth_token=os.environ['TOKEN'])
model = AutoModelForCausalLM.from_pretrained(
"EleutherAI/gpt-neox-ko-1.3b", use_auth_token=os.environ['TOKEN'],
pad_token_id=tokenizer.pad_token_id,
eos_token_id=tokenizer.eos_token_id,
bos_token_id=tokenizer.bos_token_id
)
model.eval()
pipe = pipeline('text-generation', model=model, tokenizer=tokenizer)
def predict(text):
stopping_cond = StoppingCriteriaList([tokenizer.encode('<|endoftext|>')])
with torch.no_grad():
tokens = tokenizer(text, return_tensors="pt").input_ids
gen_tokens = model.generate(tokens, do_sample=True, temperature=0.8, max_new_tokens=64)
generated = tokenizer.batch_decode(gen_tokens)[0]
return generated
# return pipe(text)[0]['generated_text']
iface = gr.Interface(
fn=predict,
inputs='text',
outputs='text',
examples=examples
)
iface.launch()
# print(generated) # print: μΈκ°„μ²˜λŸΌ μƒκ°ν•˜κ³ , ν–‰λ™ν•˜λŠ” '지λŠ₯'을 톡해 인λ₯˜κ°€ μ΄μ œκΉŒμ§€ 풀지 λͺ»ν–ˆλ˜ 문제의 해닡을 찾을 수 μžˆμ„ 것이닀. κ³Όν•™κΈ°μˆ μ΄ κ³ λ„λ‘œ λ°œλ‹¬ν•œ 21μ„ΈκΈ°λ₯Ό μ‚΄μ•„κ°ˆ 우리 μ•„μ΄λ“€μ—κ²Œ κ°€μž₯ ν•„μš”ν•œ 것은 사고λ ₯ ν›ˆλ ¨μ΄λ‹€. 사고λ ₯ ν›ˆλ ¨μ„ 톡해, 세상