| 6 | from jsonargparse import CLI |
| 7 | |
| 8 | def main( |
| 9 | checkpoint_path: Path = Path('checkpoint'), |
| 10 | prompt: str = "???", |
| 11 | max_new_tokens: int = 50, |
| 12 | repetition_penalty: float = 1.0, |
| 13 | temperature: float = 0.8, |
| 14 | tokenizer_path: Path = Path('tokenizer'), |
| 15 | ): |
| 16 | checkpoint_path = normpath(join(getcwd(), checkpoint_path)) |
| 17 | print('checkpoint_path: ', checkpoint_path) |
| 18 | checkpoint_path = Path(checkpoint_path) |
| 19 | |
| 20 | tokenizer_path = normpath(join(getcwd(), tokenizer_path)) |
| 21 | print('tokenizer_path: ', tokenizer_path) |
| 22 | tokenizer_path = Path(tokenizer_path) |
| 23 | |
| 24 | assert checkpoint_path.is_dir(), checkpoint_path |
| 25 | assert tokenizer_path.is_dir(), tokenizer_path |
| 26 | |
| 27 | |
| 28 | tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, device_map='cuda') |
| 29 | |
| 30 | prompt = f'''<|im_start|>system |
| 31 | You are a helpful, respectful, and honest assistant.<|im_end|> |
| 32 | <|im_start|>user |
| 33 | {prompt}<|im_end|> |
| 34 | <|im_start|>assistant\n''' |
| 35 | |
| 36 | input_ids = tokenizer.encode(prompt, return_tensors='pt').cuda() |
| 37 | input_ids_len = input_ids.shape[1] |
| 38 | |
| 39 | # Load the model |
| 40 | model = AutoModelForCausalLM.from_pretrained(checkpoint_path, device_map='cuda', attn_implementation="eager") |
| 41 | model.resize_token_embeddings(len(tokenizer)) |
| 42 | |
| 43 | # Generate output |
| 44 | answer_ids = model.generate( |
| 45 | input_ids, |
| 46 | max_new_tokens=max_new_tokens, |
| 47 | pad_token_id=tokenizer.pad_token_id, |
| 48 | temperature=temperature, |
| 49 | repetition_penalty=repetition_penalty |
| 50 | )[0][input_ids_len:] |
| 51 | print(tokenizer.pad_token_id) |
| 52 | # Decode the generated answer |
| 53 | answer = tokenizer.decode(answer_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) |
| 54 | print(answer) |
| 55 | |
| 56 | if __name__ == '__main__': |
| 57 | CLI(main) |