| 6 | |
| 7 | |
| 8 | def load(tokenizer_path, checkpoint_path, use_cpu=False): |
| 9 | print('Loading tokenizer ...') |
| 10 | tokenizer = AutoTokenizer.from_pretrained( |
| 11 | tokenizer_path, use_fast=False, trust_remote_code=True, padding_side='left') |
| 12 | tokenizer.add_tokens("[USER]") |
| 13 | tokenizer.add_tokens("[BOT]") |
| 14 | tokenizer.add_tokens("[SEP]") |
| 15 | |
| 16 | print('Loading model ...') |
| 17 | config = AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=True) |
| 18 | if use_cpu: |
| 19 | device_map = "cpu" |
| 20 | else: |
| 21 | device_map = "balanced_low_0" |
| 22 | |
| 23 | model = AutoModelForCausalLM.from_pretrained( |
| 24 | checkpoint_path, config=config, device_map=device_map, torch_dtype=torch.bfloat16, trust_remote_code=True) |
| 25 | model.generation_config = GenerationConfig.from_pretrained( |
| 26 | checkpoint_path, trust_remote_code=True) |
| 27 | |
| 28 | model.generation_config.do_sample = True |
| 29 | model.eval() |
| 30 | |
| 31 | if tokenizer.pad_token_id is None: |
| 32 | if tokenizer.eos_token_id is not None: |
| 33 | tokenizer.pad_token_id = tokenizer.eos_token_id |
| 34 | else: |
| 35 | tokenizer.pad_token_id = 0 |
| 36 | |
| 37 | return model, tokenizer |
| 38 | |
| 39 | def special_encode(prompt, tokenizer): |
| 40 | raw_str = "[USER]%s[SEP][BOT]" % prompt.strip().replace("\r", "") |