Moo
/

professional_comment

Text Generation

Inference Endpoints

Model card Files Files and versions Community

professional_comment / inference.py

Moo's picture

Moo

Upload inference.py

53f9185 almost 2 years ago

raw history blame contribute delete

No virus

2.33 kB

	import torch
	import transformers
	from transformers import AutoTokenizer, GPTJForCausalLM

	def get_sent(sent:str) -> str:
	input_text = '[BOS]' + sent + '[EOS][BOS]'
	input_length = len(tokenizer.encode(input_text))
	max_length = 786
	input_ids = torch.tensor(tokenizer.encode(input_text)).unsqueeze(0).to('cuda')
	output_sentence = model.generate(
	input_ids,
	do_sample=True,
	max_length=int(max_length),
	num_return_sequences=1,
	no_repeat_ngram_size=4,
	num_beams=5,
	early_stopping=True
	)
	generated_sequence = output_sentence[0].tolist()[input_length:]
	decoded_sent = tokenizer.decode(generated_sequence, skip_special_tokens=False).strip()
	return decoded_sent


	if __name__ == "__main__":
	tokenizer = AutoTokenizer.from_pretrained('kakaobrain/kogpt', revision="KoGPT6B-ryan1.5b")
	model = GPTJForCausalLM.from_pretrained('./', torch_dtype=torch.float16)
	model.cuda()

	text = """
	번아웃이 온것을 인정하기 싫어요.
	2년전부터 우울증에 시달렸습니다. 2년동안 저를 알아가려고 애쓰고 동시에 저의 미래를 위해 끊임없이 자격증 공부와 학업을 병행해 왔어요.
	하지만 불안정한 심리 상태로 인하여 제 꿈과 수많은 기회들을 잡지 못하고 도망쳐버렸습니다. 어째서일까요 왜 저는 남들보다 버티는 힘이 약한걸까요.
	다시 힘을내서 도전도 해보고 여러 일을 해보며 모든 에너지를 쏟아부었더니 어느날 한글이 잘 파악이 안되고 간단한 결정을 내리는 것이 어려워 결국 직장을 또 그만두게 되었습니다. 나약하다고 하기에는 보통사람들이 버티기 어려워하는 직종을 하고 있거든요.
	너무 과도하게 힘을 내서 열심히 한 탓일까요? 이제는 정말 배터리가 1%도 남아있지 않는 것 같아요. 근데 쉬려니까 불안해지고 부모님께 미안해요 아프지만 않았어도 나도 자리 잘 잡아서 잘 지냈을 텐데요.. 앞으로 저는 그냥 휴식을 취하는 것이 맞을까요? 저는 어떤 상태에 있는 걸까요?
	"""

	text = text.replace('\n','')
	result = get_sent(text)
	result = result.replace('사우님', '마카님')
	print('로니형:', result)