(self, cfg: Config, *, parent: Optional[Module])
| 796 | transducer: Transducer.Config = Transducer.default_config() |
| 797 | |
| 798 | def __init__(self, cfg: Config, *, parent: Optional[Module]): |
| 799 | super().__init__(cfg, parent=parent) |
| 800 | cfg = self.config |
| 801 | if cfg.eos_id == cfg.blank_id: |
| 802 | raise ValueError( |
| 803 | "eos_id and blank_id should be different for the transducer model, " |
| 804 | f"but got eos_id = blank_id = {cfg.blank_id}." |
| 805 | ) |
| 806 | self.vlog( |
| 807 | 3, |
| 808 | ( |
| 809 | f"am_dim={cfg.input_dim}, lm_dim={cfg.lm_dim}, joint_dim={cfg.joint_dim}, " |
| 810 | f"vocab_size={cfg.vocab_size}." |
| 811 | ), |
| 812 | ) |
| 813 | # In most common cases, am_data and lm_data are summed together after the projection, thus |
| 814 | # we only keep one bias in the two projections. |
| 815 | self._add_child( |
| 816 | "am_proj", cfg.am_proj.set(input_dim=cfg.input_dim, output_dim=cfg.joint_dim, bias=True) |
| 817 | ) |
| 818 | self._add_child( |
| 819 | "lm_proj", cfg.lm_proj.set(input_dim=cfg.lm_dim, output_dim=cfg.joint_dim, bias=False) |
| 820 | ) |
| 821 | self._add_child( |
| 822 | "prediction_network", |
| 823 | cfg.prediction_network.set( |
| 824 | vocab_size=cfg.vocab_size, |
| 825 | output_dim=cfg.lm_dim, |
| 826 | ), |
| 827 | ) |
| 828 | transducer_cfg = cfg.transducer.set(input_dim=cfg.joint_dim, vocab_size=cfg.vocab_size) |
| 829 | transducer_cfg.logits_to_log_probs.blank_id = cfg.blank_id |
| 830 | self._add_child("transducer", transducer_cfg) |
| 831 | |
| 832 | def forward(self, input_batch: Nested[Tensor]) -> tuple[Tensor, Nested[Tensor]]: |
| 833 | """Computes the transducer loss. |
nothing calls this directly
no test coverage detected