Encode source label sequences. Args: x: Encoder input sequences. (B, L) Returns: out: Encoder output sequences. (B, U, D)
(self, x: torch.Tensor, x_len)
| 96 | return self._output_size |
| 97 | |
| 98 | def forward(self, x: torch.Tensor, x_len) -> torch.Tensor: |
| 99 | """Encode source label sequences. |
| 100 | |
| 101 | Args: |
| 102 | x: Encoder input sequences. (B, L) |
| 103 | |
| 104 | Returns: |
| 105 | out: Encoder output sequences. (B, U, D) |
| 106 | |
| 107 | """ |
| 108 | _, length, _ = x.size() |
| 109 | |
| 110 | assert ( |
| 111 | length <= self.context_size * self.subsampling_factor |
| 112 | ), "Context size is too short for current length: %d versus %d" % ( |
| 113 | length, |
| 114 | self.context_size * self.subsampling_factor, |
| 115 | ) |
| 116 | mask = make_source_mask(x_len).to(x.device) |
| 117 | x, mask = self.embed(x, mask, None) |
| 118 | x = self.embed_norm(x) |
| 119 | olens = mask.eq(0).sum(1) |
| 120 | |
| 121 | if self.training: |
| 122 | for block in self.rwkv_blocks: |
| 123 | x, _ = block(x) |
| 124 | else: |
| 125 | x = self.rwkv_infer(x) |
| 126 | |
| 127 | x = self.final_norm(x) |
| 128 | |
| 129 | if self.time_reduction_factor > 1: |
| 130 | x = x[:, :: self.time_reduction_factor, :] |
| 131 | olens = torch.floor_divide(olens - 1, self.time_reduction_factor) + 1 |
| 132 | |
| 133 | return x, olens, None |
| 134 | |
| 135 | def rwkv_infer(self, xs_pad): |
| 136 |
nothing calls this directly
no test coverage detected