github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[chatgpt] fix inference demo loading bug (#2969) 

* [chatgpt] fix inference demo loading bug

* polish
pull/2974/head
BlueRum 2023-03-02 16:18:33 +08:00 committed by GitHub
parent bbf9c827c3
commit 82149e9d1b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 12 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
applications/ChatGPT/examples/inference.py
View File

@ -9,30 +9,34 @@ from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
def eval(args):
# configure model
if args.model == 'gpt2':
model = GPTActor(pretrained=args.pretrain).to(torch.cuda.current_device())
actor = GPTActor().to(torch.cuda.current_device())
elif args.model == 'bloom':
model = BLOOMActor(pretrained=args.pretrain).to(torch.cuda.current_device())
actor = BLOOMActor().to(torch.cuda.current_device())
elif args.model == 'opt':
model = OPTActor(pretrained=args.pretrain).to(torch.cuda.current_device())
actor = OPTActor().to(torch.cuda.current_device())
else:
raise ValueError(f'Unsupported model "{args.model}"')
state_dict = torch.load(args.pretrain)
actor.model.load_state_dict(state_dict)
# configure tokenizer
if args.model == 'gpt2':
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
tokenizer.pad_token = tokenizer.eos_token
elif args.model == 'bloom':
tokenizer = AutoTokenizer.from_pretrained(args.pretrain)
tokenizer = AutoTokenizer.from_pretrained('bigscience/bloom-560m')
tokenizer.pad_token = tokenizer.eos_token
elif args.model == 'opt':
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m')
else:
raise ValueError(f'Unsupported model "{args.model}"')
model.eval()
actor.eval()
input = args.input
input_ids = tokenizer.encode(input, return_tensors='pt').to(torch.cuda.current_device())
outputs = model.generate(input_ids,
outputs = actor.generate(input_ids,
max_length=args.max_length,
do_sample=True,
top_k=50,
@ -46,7 +50,7 @@ if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt'])
parser.add_argument('--pretrain', type=str, default=None)
parser.add_argument('--input', type=str, default='Q: How are you ? A:')
parser.add_argument('--input', type=str, default='Question: How are you ? Answer:')
parser.add_argument('--max_length', type=int, default=100)
args = parser.parse_args()
eval(args)