mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-03 04:39:43 +00:00
da4f7b855f
* style: rename replay buffer Experience replay is typically for off policy algorithms. Use this name in PPO maybe misleading. * fix: fix wrong zero2 default arg * test: update experience tests * style: rename zero_pad fn * fix: defer init in CycledDataLoader * test: add benchmark test * style: rename internal fn of generation * style: rename internal fn of lora * fix: remove unused loss fn * fix: remove unused utils fn * refactor: remove generate_with_actor fn * fix: fix type annotation * test: add models tests * fix: skip llama due to long execution time * style: modify dataset * style: apply formatter * perf: update reward dataset * fix: fix wrong IGNORE_INDEX in sft dataset * fix: remove DataCollatorForSupervisedDataset * test: add dataset tests * style: apply formatter * style: rename test_ci to test_train * feat: add llama in inference * test: add inference tests * test: change test scripts directory * fix: update ci * fix: fix typo * fix: skip llama due to oom * fix: fix file mod * style: apply formatter * refactor: remove duplicated llama_gptq * style: apply formatter * to: update rm test * feat: add tokenizer arg * feat: add download model script * test: update train tests * fix: modify gemini load and save pretrained * test: update checkpoint io test * to: modify nproc_per_node * fix: do not remove existing dir * fix: modify save path * test: add random choice * fix: fix sft path * fix: enlarge nproc_per_node to avoid oom * fix: add num_retry * fix: make lora config of rm and critic consistent * fix: add warning about lora weights * fix: skip some gpt2 tests * fix: remove grad ckpt in rm and critic due to errors * refactor: directly use Actor in train_sft * test: add more arguments * fix: disable grad ckpt when using lora * fix: fix save_pretrained and related tests * test: enable zero2 tests * revert: remove useless fn * style: polish code * test: modify test args
73 lines
2.8 KiB
Python
73 lines
2.8 KiB
Python
import argparse
|
|
|
|
import torch
|
|
from coati.models.bloom import BLOOMActor
|
|
from coati.models.generation import generate
|
|
from coati.models.gpt import GPTActor
|
|
from coati.models.llama import LlamaActor
|
|
from coati.models.opt import OPTActor
|
|
from transformers import AutoTokenizer, BloomTokenizerFast, GPT2Tokenizer, LlamaTokenizer
|
|
|
|
|
|
def eval(args):
|
|
# configure model
|
|
if args.model == 'gpt2':
|
|
actor = GPTActor(pretrained=args.pretrain)
|
|
elif args.model == 'bloom':
|
|
actor = BLOOMActor(pretrained=args.pretrain)
|
|
elif args.model == 'opt':
|
|
actor = OPTActor(pretrained=args.pretrain)
|
|
elif args.model == 'llama':
|
|
actor = LlamaActor(pretrained=args.pretrain)
|
|
else:
|
|
raise ValueError(f'Unsupported model "{args.model}"')
|
|
|
|
actor.to(torch.cuda.current_device())
|
|
if args.model_path is not None:
|
|
state_dict = torch.load(args.model_path)
|
|
actor.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 = BloomTokenizerFast.from_pretrained('bigscience/bloom-560m')
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
elif args.model == 'opt':
|
|
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
elif args.model == 'llama':
|
|
tokenizer = LlamaTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
|
|
tokenizer.eos_token = '<\s>'
|
|
tokenizer.pad_token = tokenizer.unk_token
|
|
else:
|
|
raise ValueError(f'Unsupported model "{args.model}"')
|
|
|
|
actor.eval()
|
|
input_ids = tokenizer.encode(args.input,
|
|
return_tensors='pt')\
|
|
.to(torch.cuda.current_device())
|
|
outputs = generate(actor,
|
|
input_ids,
|
|
max_length=args.max_length,
|
|
do_sample=True,
|
|
top_k=50,
|
|
top_p=0.95,
|
|
num_return_sequences=1)
|
|
output = tokenizer.batch_decode(outputs[0],
|
|
skip_special_tokens=True)
|
|
print(f"[Output]: {''.join(output)}")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt', 'llama'])
|
|
# We suggest to use the pretrained model from HuggingFace, use pretrain to configure model
|
|
parser.add_argument('--pretrain', type=str, default=None)
|
|
parser.add_argument('--model_path', type=str, default=None)
|
|
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)
|