src/model/extract_embeddings.py

import torch

def extract_embeddings(model, padded_sequences, batch_size=64):
    device = next(model.parameters()).device

    model.eval()

    data = torch.tensor(padded_sequences, dtype=torch.long)
    dataset = torch.utils.data.TensorDataset(data)
    loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)

    all_embeddings = []

    with torch.no_grad():
        for batch in loader:
            batch = batch[0].to(device)

            emb = model.get_embeddings(batch)  # (B, D)
            all_embeddings.append(emb.cpu())

    return torch.cat(all_embeddings, dim=0)  # (N, D)