| 34 | |
| 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): |
| 46 | """Run one forward step and return next sentence + its embedding.""" |
| 47 | if embedded_sents.ndim == 2: |
| 48 | embedded_sents = embedded_sents.unsqueeze(0) # [1, T, dim] |
| 49 | |
| 50 | proj = self.forward_proj(embedded_sents) |
| 51 | out = self.llama_model(inputs_embeds=proj, output_hidden_states=True) |
| 52 | hidden = out.hidden_states[-1] # [1, T, hidden] |
| 53 | final_hidden = hidden[0, -1, :] |
| 54 | reversed_emb = self.reverse_proj(final_hidden.unsqueeze(0)) |
| 55 | out_texts = self.sonar_decoder.predict( |
| 56 | reversed_emb, target_lang="eng_Latn", max_seq_len=256 |
| 57 | ) |
| 58 | return out_texts[0], reversed_emb |
| 59 | |
| 60 | |
| 61 | def load_checkpoint_for_inference(ckpt_path, model, device="cpu"): |