Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4tokenizer = AutoTokenizer.from_pretrained(
5 "CheonggyeMountain-Sherpa/kogpt-trinity-punct-wrapper",
6 revision="punct_wrapper-related_words-overfit", # or punct_wrapper-related_words-minevalloss
7 bos_token="<s>",
8 eos_token="</s>",
9 unk_token="<unk>",
10 pad_token="<pad>",
11 mask_token="<mask>",
12)
13model = AutoModelForCausalLM.from_pretrained(
14 "CheonggyeMountain-Sherpa/kogpt-trinity-punct-wrapper",
15 revision="punct_wrapper-related_words-overfit", # or punct_wrapper-related_words-minevalloss
16 pad_token_id=tokenizer.eos_token_id,
17).to(device="cuda")
18model.eval()
19
20prompt = "석양이 보이는 경치"
21wrapped_prompt = f"@{prompt}@<usr>\n"
22with torch.no_grad():
23 tokens = tokenizer.encode(wrapped_prompt, return_tensors="pt").to(device="cuda")
24 gen_tokens = model.generate(
25 tokens,
26 max_length=64,
27 repetition_penalty=2.0,
28 pad_token_id=tokenizer.pad_token_id,
29 eos_token_id=tokenizer.eos_token_id,
30 bos_token_id=tokenizer.bos_token_id,
31 top_k=16,
32 top_p=0.8,
33 )
34 generated = tokenizer.decode(gen_tokens[0][len(tokens[0]):])
35
36print(generated)
37# 해가 지고 있을 무렵
38# 나는 석양을 보러 간다
39# 붉은 하늘과 하얀 구름이 나를 반겨줄 것 같아서리
40# 하지만 내가 본 해는 저물어만 가고
41# 구름마저 자취를 감춘 어둠만이 남아있을 뿐이네
42# 내가 탄 배는 보이지도 않고