| 324 | elif args.patch_tokenizer_type == 'MistralTokenizer': |
| 325 | from megatron.core.datasets.megatron_tokenizer import MegatronTokenizer |
| 326 | class _MistralTokenizer(MegatronTokenizer): |
| 327 | def __init__(self, tokenizer_path, extra_vocab_size): |
| 328 | super().__init__(tokenizer_path) |
| 329 | self.tokenizer = AutoTokenizer.from_pretrained( |
| 330 | tokenizer_path, |
| 331 | padding_side="right", |
| 332 | use_fast=False, |
| 333 | trust_remote_code=True |
| 334 | ) |
| 335 | self.extra_vocab_size = extra_vocab_size |
| 336 | |
| 337 | def __call__(self, text, return_tensors=None, |
| 338 | padding=None, max_length=None, truncation=None, add_special_tokens=None): |
| 339 | |
| 340 | return self.tokenizer(text, return_tensors=return_tensors, padding=padding, |
| 341 | max_length=max_length, truncation=truncation, add_special_tokens=add_special_tokens) |
| 342 | |
| 343 | @property |
| 344 | def vocab_size(self): |
| 345 | return self.tokenizer.vocab_size + self.extra_vocab_size |
| 346 | |
| 347 | @property |
| 348 | def vocab(self): |
| 349 | return self.tokenizer.encoder |
| 350 | |
| 351 | @property |
| 352 | def inv_vocab(self): |
| 353 | return self.tokenizer.decoder |
| 354 | |
| 355 | def tokenize(self, text): |
| 356 | return self.tokenizer.encode(text) |
| 357 | |
| 358 | def detokenize(self, token_ids): |
| 359 | return self.tokenizer.decode(token_ids) |
| 360 | |
| 361 | @property |
| 362 | def eod(self): |
| 363 | return self.tokenizer.eos_token_id |
| 364 | |
| 365 | @property |
| 366 | def eos_token(self): |
| 367 | return self.tokenizer.eos_token |
| 368 | |
| 369 | @property |
| 370 | def pad_token_id(self): |
| 371 | return self.tokenizer.pad_token_id |
| 372 | |
| 373 | @property |
| 374 | def eos_token_id(self): |
| 375 | return self.tokenizer.eos_token_id |
| 376 | |
| 377 | tokenizer = _MistralTokenizer(args.load, args.extra_vocab_size) |
| 378 | args.padded_vocab_size = tokenizer.vocab_size + args.extra_vocab_size |