(model_name: str, is_hf: bool=False)
| 95 | |
| 96 | |
| 97 | def load_tokenizer(model_name: str, is_hf: bool=False) -> nn.Module: |
| 98 | if not is_hf: |
| 99 | tokenizer = AutoTokenizer.from_pretrained("gpt2") |
| 100 | tokenizer.model_max_length = 2048 |
| 101 | else: |
| 102 | if "mamba" in model_name or "mpt" in model_name: |
| 103 | tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") |
| 104 | else: |
| 105 | tokenizer = AutoTokenizer.from_pretrained(model_name) |
| 106 | tokenizer.pad_token = tokenizer.eos_token |
| 107 | tokenizer.pad_token_id = tokenizer.eos_token_id |
| 108 | return tokenizer |
| 109 | |
| 110 |