(self, llama_model, forward_proj, reverse_proj, sonar_decoder)
| 35 | |
| 36 | class SonarInferenceWrapper(nn.Module): |
| 37 | def __init__(self, llama_model, forward_proj, reverse_proj, sonar_decoder): |
| 38 | super().__init__() |
| 39 | self.llama_model = llama_model |
| 40 | self.forward_proj = forward_proj |
| 41 | self.reverse_proj = reverse_proj |
| 42 | self.sonar_decoder = sonar_decoder |
| 43 | |
| 44 | @torch.no_grad() |
| 45 | def inference_step(self, embedded_sents: torch.Tensor): |