(
self,
text_len,
dtype=torch.bfloat16,
device=torch.cuda.current_device(),
checkpoint_path=None,
tokenizer_path=None,
shard_fn=None,
)
| 541 | class T5EncoderModel: |
| 542 | |
| 543 | def __init__( |
| 544 | self, |
| 545 | text_len, |
| 546 | dtype=torch.bfloat16, |
| 547 | device=torch.cuda.current_device(), |
| 548 | checkpoint_path=None, |
| 549 | tokenizer_path=None, |
| 550 | shard_fn=None, |
| 551 | ): |
| 552 | self.text_len = text_len |
| 553 | self.dtype = dtype |
| 554 | self.device = device |
| 555 | self.checkpoint_path = checkpoint_path |
| 556 | self.tokenizer_path = tokenizer_path |
| 557 | |
| 558 | # init model |
| 559 | with init_empty_weights(): |
| 560 | model = umt5_xxl( |
| 561 | encoder_only=True, |
| 562 | return_tokenizer=False, |
| 563 | dtype=dtype, |
| 564 | device=device).eval().requires_grad_(False) |
| 565 | |
| 566 | if checkpoint_path.endswith('.safetensors'): |
| 567 | state_dict = load_file(checkpoint_path, device='cpu') |
| 568 | state_dict = umt5_keys_mapping(state_dict) |
| 569 | else: |
| 570 | state_dict = torch.load(checkpoint_path, map_location='cpu') |
| 571 | |
| 572 | model.load_state_dict(state_dict, assign=True) |
| 573 | self.model = model |
| 574 | if shard_fn is not None: |
| 575 | self.model = shard_fn(self.model, sync_module_states=False) |
| 576 | else: |
| 577 | self.model.to(self.device) |
| 578 | # init tokenizer |
| 579 | self.tokenizer = HuggingfaceTokenizer( |
| 580 | name=tokenizer_path, seq_len=text_len, clean='whitespace') |
| 581 | |
| 582 | def __call__(self, texts, device): |
| 583 | ids, mask = self.tokenizer( |
nothing calls this directly
no test coverage detected