import torch\nimport torch.nn.functional as F\n\nclass ContrastiveDomainAdapter(nn.Module):\n def __init__(self, encoder, temp=0.07):\n super().__init__()\n self.encoder = encoder\n self.projection = nn.Linear(768, 768)\n self.temperature = temp\n\n def contrastive_loss(self, z, labels):\n z = F.normalize(z, dim=-1)\n sim = torch.mm(z, z.t()) / self.temperature\n mask = torch.eq(labels.unsqueeze(0), labels.unsqueeze(1)).float()\n mask.fill_diagonal_(0)\n pos = (sim * mask).sum(dim=-1)\n neg = torch.logsumexp(sim * (1 - mask), dim=-1)\n return (-pos + neg).mean()