Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
2import torch
3import time
4import random
5import numpy as np
6
7#
8# Fix seed
9#
10seed = 42
11
12random.seed(seed)
13# Numpy
14np.random.seed(seed)
15# Pytorch
16torch.manual_seed(seed)
17torch.cuda.manual_seed(seed)
18torch.backends.cudnn.deterministic = True
19torch.use_deterministic_algorithms = True
20
21torch.set_default_dtype(torch.bfloat16)
22
23
24
25model_id = "aerner/lm-v1"
26
27
28text = """### Instruction:
29東京駅について説明してください。
30
31
32### Context:
33
34
35
36### Answer:
37"""
38
39with torch.no_grad():
40 tokenizer = AutoTokenizer.from_pretrained(model_id)
41 tokenized_input = tokenizer(text, return_tensors="pt").to('cuda')
42
43 model = AutoModelForCausalLM.from_pretrained(
44 model_id, device_map="auto", torch_dtype=torch.bfloat16)
45
46 generation_config = GenerationConfig(
47 max_new_tokens=256,
48 min_new_tokens=1,
49 early_stopping=True,
50 do_sample=True,
51
52 num_beams=8,
53 temperature=1.0,
54 top_p=0.6,
55 penalty_alpha=0.4,
56 no_repeat_ngram_size=4,
57 repetition_penalty=1.4,
58
59 remove_invalid_values=True,
60 num_return_sequences=1,
61 )
62
63 start = time.time()
64
65 generation_output = model.generate(
66 input_ids=tokenized_input['input_ids'],
67 generation_config=generation_config,
68 return_dict_in_generate=True,
69 output_scores=True,
70 )
71
72 for s in generation_output.sequences:
73 output = tokenizer.decode(s)
74 print(output)
75
76 print(time.time() - start)
77