Method forward

train_sonarllm.py:168–226

(self, embeddings_1024, texts, seq_lens)

            p.requires_grad = False

    def forward(self, embeddings_1024, texts, seq_lens):
        device = embeddings_1024.device
        B, T, _ = embeddings_1024.shape

        embs_projected = self.forward_proj(embeddings_1024)
        llama_out = self.llama_model(
            inputs_embeds=embs_projected,
            output_hidden_states=True
        )
        last_hidden = llama_out.hidden_states[-1]

        pred_hidden_list = []
        ref_texts_list = []
        for b in range(B):
            seqlen = seq_lens[b]
            for k in range(1, seqlen):
                pred_hidden_list.append(last_hidden[b, k - 1, :])
                ref_texts_list.append(texts[b][k])

        if len(pred_hidden_list) == 0:
            return torch.tensor(0.0, device=device)

        pred_hidden_batch = torch.stack(pred_hidden_list, dim=0)
        pred_emb_1024 = self.reverse_proj(pred_hidden_batch)

        with torch.no_grad():
            target_text_encoder = self.sonar_decoder.tokenizer.create_encoder(
                task="translation", lang="eng_Latn", mode="target", device=device
            )
        encoded_texts = [target_text_encoder(t) for t in ref_texts_list]
        lengths = [et.size(0) for et in encoded_texts]
        max_len = min(max(lengths), 256)

        pad_idx = self.sonar_decoder.tokenizer.vocab_info.pad_idx
        dec_ids = torch.full((len(encoded_texts), max_len), pad_idx, dtype=torch.long, device=device)
        labels = torch.full((len(encoded_texts), max_len), pad_idx, dtype=torch.long, device=device)
        for i, et in enumerate(encoded_texts):
            dec_ids[i, : min(len(et), max_len)] = et[:max_len]
            et = torch.cat([et[1:], torch.tensor([3]).to(device)])
            labels[i, : min(len(et), max_len)] = et[:max_len]

        enc_output = pred_emb_1024.unsqueeze(1)
        dec_out, dec_pad_mask = self.sonar_decoder.model.decode(
            seqs=dec_ids,
            padding_mask=None,
            encoder_output=enc_output,
            encoder_padding_mask=None
        )
        final_out = self.sonar_decoder.model.project(dec_out, dec_pad_mask)
        logits = final_out.logits

        vocab_size = logits.size(-1)
        logits_2d = logits.view(-1, vocab_size)
        labels_1d = labels.view(-1)

        ce_fn = nn.CrossEntropyLoss(ignore_index=pad_idx, reduction="mean")
        total_ce = ce_fn(logits_2d, labels_1d)