Freeze token embeddings and positional embeddings for bart, just token embeddings for t5.
(model)
| 548 | |
| 549 | |
| 550 | def freeze_embeds(model): |
| 551 | """Freeze token embeddings and positional embeddings for bart, just token embeddings for t5.""" |
| 552 | model_type = model.config.model_type |
| 553 | |
| 554 | if model_type == "t5": |
| 555 | freeze_params(model.shared) |
| 556 | for d in [model.encoder, model.decoder]: |
| 557 | freeze_params(d.embed_tokens) |
| 558 | elif model_type == "fsmt": |
| 559 | for d in [model.model.encoder, model.model.decoder]: |
| 560 | freeze_params(d.embed_positions) |
| 561 | freeze_params(d.embed_tokens) |
| 562 | else: |
| 563 | freeze_params(model.model.shared) |
| 564 | for d in [model.model.encoder, model.model.decoder]: |
| 565 | freeze_params(d.embed_positions) |
| 566 | freeze_params(d.embed_tokens) |
| 567 | |
| 568 | |
| 569 | def grad_status(model: nn.Module) -> Iterable: |
nothing calls this directly
no test coverage detected