(self, weights, from_pruned=False)
| 530 | return model |
| 531 | |
| 532 | def load(self, weights, from_pruned=False): |
| 533 | required_names = set() |
| 534 | for name, param in self.named_parameters(): |
| 535 | if self.use_pretrained_pos_emb and 'pos_embed' in name: |
| 536 | required_names.add(name) |
| 537 | continue |
| 538 | if param.is_inited(): |
| 539 | continue |
| 540 | if name not in weights: |
| 541 | # Exemption for embedding sharing |
| 542 | if name.endswith('lm_head.weight') and any( |
| 543 | k.endswith('vocab_embedding.weight') |
| 544 | for k in weights.keys()): |
| 545 | continue |
| 546 | if name.endswith('lm_head.per_channel_scale') and any( |
| 547 | k.endswith('vocab_embedding.per_channel_scale') |
| 548 | for k in weights.keys()): |
| 549 | continue |
| 550 | required_names.add(name) |
| 551 | |
| 552 | provided_names = set(weights.keys()) |
| 553 | if not required_names.issubset(provided_names): |
| 554 | raise RuntimeError( |
| 555 | f"Required but not provided tensors:{required_names.difference(provided_names)}" |
| 556 | ) |
| 557 | if not provided_names.issubset(required_names): |
| 558 | logger.warning( |
| 559 | f"Provided but not required tensors: {provided_names.difference(required_names)}" |
| 560 | ) |
| 561 | |
| 562 | for name, param in self.named_parameters(): |
| 563 | if name in provided_names: |
| 564 | if not from_pruned: |
| 565 | try: |
| 566 | param.value = weights[name] |
| 567 | except Exception as e: |
| 568 | raise RuntimeError( |
| 569 | f"Encounter error '{e}' for parameter '{name}'") |
| 570 | else: |
| 571 | param.set_value_or_dummy(weights[name]) |
| 572 | |
| 573 | def enable_forward_chunking(self, |
| 574 | chunk_size: Optional[int] = None, |
no test coverage detected