Views
No views yet
pip install transformers sentencepiece accelerate1from transformers import GPTJForCausalLM, AlbertTokenizer
2import torch
3
4tokenizer = AlbertTokenizer.from_pretrained('AIBunCho/japanese-novel-gpt-j-6b', keep_accents=True, remove_space=False)
5
6model = GPTJForCausalLM.from_pretrained("AIBunCho/japanese-novel-gpt-j-6b", torch_dtype=torch.float16, low_cpu_mem_usage=True)
7
8model.half()
9model.eval()
10
11if torch.cuda.is_available():
12 model = model.to("cuda")
13
14prompt = """
15わたくしといふ現象は
16""".strip()
17
18input_ids = tokenizer.encode(
19 prompt,
20 add_special_tokens=False,
21 return_tensors="pt"
22).cuda()
23
24# this is for reproducibility.
25# feel free to change to get different result
26seed = 27
27torch.manual_seed(seed)
28
29tokens = model.generate(
30 input_ids.to(device=model.device),
31 max_new_tokens=32,
32 temperature=0.6,
33 top_p=0.9,
34 repetition_penalty=1.2,
35 do_sample=True,
36 pad_token_id=tokenizer.pad_token_id,
37 bos_token_id=tokenizer.bos_token_id,
38 eos_token_id=tokenizer.eos_token_id
39)
40
41out = tokenizer.decode(tokens[0], skip_special_tokens=True)
42print(out)
43"""わたくしといふ現象は、その因果律を断ち切ることができるのです。"""
44