Spaces:

JYYong
/

test_space

Runtime error

App Files Files Community

test_space / app.py

JYYong

maybe complete

4cb51c5 almost 2 years ago

raw history blame contribute delete

No virus

5.48 kB

	import gradio as gr
	from transformers import AutoTokenizer, AutoModelForCausalLM
	import warnings


	class Chatbot():
	def __init__(self):
	self.tokenizer = AutoTokenizer.from_pretrained('kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b')
	special_tokens_dict = {'additional_special_tokens': ['<sep>', '<eos>', '<sos>', '#@이름#', '#@계정#', '#@신원#', '#@전번#', '#@금융#', '#@번호#', '#@주소#', '#@소속#', '#@기타#']}
	num_added_toks = self.tokenizer.add_special_tokens(special_tokens_dict)

	self.model = AutoModelForCausalLM.from_pretrained("/workspace/test_trainer/checkpoint-10000")
	self.model.resize_token_embeddings(len(self.tokenizer))
	self.model = self.model.cuda()

	self.info = None
	self.talk = []

	def initialize(self, topic, bot_addr, bot_age, bot_sex, my_addr, my_age, my_sex):
	def encode(age):
	if age < 20:
	age = "20대 미만"
	elif age >= 70:
	age = "70대 이상"
	else:
	age = str(age // 10 * 10) + "대"
	return age
	bot_age = encode(bot_age)
	my_age = encode(my_age)
	self.info = f"일상 대화 {topic}<sep>P01:{my_addr} {my_age} {my_sex}<sep>P02:{bot_addr} {bot_age} {bot_sex}<sep>"
	return self.info_check()

	def info_check(self):
	return self.info.replace('<sep>', '\n').replace('P01', '당신').replace('P02', '챗봇')

	def reset_talk(self):
	self.talk = []

	def test(self, myinp):
	state = None
	inp = "P01<sos>" + myinp + "<eos>"
	self.talk.append(inp)
	self.talk.append("P02<sos>")

	while True:
	now_inp = self.info + "".join(self.talk)
	inputs = self.tokenizer(now_inp, max_length=1024, truncation='longest_first', return_tensors='pt')
	seq_len = inputs.input_ids.size(1)
	if seq_len > 512 * 0.8:
	state = f"<주의> 현재 대화 길이가 곧 최대 길이에 도달합니다. ({seq_len} / 512)"

	if seq_len >= 512:
	state = "<주의> 대화 길이가 너무 길어졌기 때문에, 이후 대화는 맨 앞의 발화를 조금씩 지우면서 진행됩니다."
	talk = talk[1:]
	else:
	break

	out = self.model.generate(
	inputs=inputs.input_ids.cuda(),
	attention_mask=inputs.attention_mask.cuda(),
	max_length=512,
	do_sample=True,
	pad_token_id=self.tokenizer.pad_token_id,
	eos_token_id=self.tokenizer.encode('<eos>')[0]
	)
	out = self.tokenizer.batch_decode(out)
	real_out = out[0][len(now_inp):-5]
	self.talk[-1] += out[0][len(now_inp):]
	return [(self.talk[i][8:-5], self.talk[i+1][8:-5]) for i in range(0, len(self.talk)-1, 2)]


	if __name__ == "__main__":
	warnings.filterwarnings("ignore")

	chatbot = Chatbot()
	demo = gr.Blocks()

	with demo:
	gr.Markdown("# <center>MINDs Lab Brain's Fast Neural Chit-Chatbot</center>")
	with gr.Row():
	with gr.Column():
	topic = gr.Radio(label="Topic", choices=['여가 생활', '시사/교육', '미용과 건강', '식음료', '상거래(쇼핑)', '일과 직업', '주거와 생활', '개인 및 관계', '행사'])
	with gr.Column():
	gr.Markdown(f"Bot's persona")
	bot_addr = gr.Dropdown(label="지역", choices=['서울특별시', '경기도', '부산광역시', '대전광역시', '광주광역시', '울산광역시', '경상남도', '인천광역시', '충청북도', '제주도', '강원도', '충청남도', '전라북도', '대구광역시', '전라남도', '경상북도', '세종특별자치시', '기타'])
	bot_age = gr.Slider(label="나이", minimum=10, maximum=80, value=45, step=1)
	bot_sex = gr.Radio(label="성별", choices=["남성", "여성"])
	with gr.Column():
	gr.Markdown(f"Your persona")
	my_addr = gr.Dropdown(label="지역", choices=['서울특별시', '경기도', '부산광역시', '대전광역시', '광주광역시', '울산광역시', '경상남도', '인천광역시', '충청북도', '제주도', '강원도', '충청남도', '전라북도', '대구광역시', '전라남도', '경상북도', '세종특별자치시', '기타'])
	my_age = gr.Slider(label="나이", minimum=10, maximum=80, value=45, step=1)
	my_sex = gr.Radio(label="성별", choices=["남성", "여성"])
	with gr.Row():
	btn = gr.Button(label="적용")
	state = gr.Textbox(label="상태")
	btn.click(
	fn=chatbot.initialize,
	inputs=[topic, bot_addr, bot_age, bot_sex, my_addr, my_age, my_sex],
	outputs=state
	)

	with gr.Column():
	screen = gr.Chatbot(label="익명의 상대")
	with gr.Row():
	speak = gr.Textbox(label="입력창")
	btn = gr.Button(label="Talk")
	btn.click(
	fn=chatbot.test,
	inputs=speak,
	outputs=screen
	)
	demo.launch(share=True)