model.py

import torch
from torch import nn
import math
import pytorch_lightning as pl


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)


class Vector_SIM_PRF(nn.Module):
    def __init__(self, ninp, nhead, nhid, nlayers, dropout):
        super(Vector_SIM_PRF, self).__init__()
        from torch.nn import TransformerEncoder, TransformerEncoderLayer
        self.model_type = 'Transformer'
        self.pos_encoder = PositionalEncoding(ninp, dropout)
        encoder_layers = TransformerEncoderLayer(ninp, nhead, nhid, dropout)
        self.norm = nn.LayerNorm(ninp, eps=1e-05)
        self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers, norm=self.norm)
        self.ninp = ninp

    def forward(self, src):
        src = src * math.sqrt(self.ninp)
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src)
        return output


class Vector_PRF_Lightning_CLS_SIM(pl.LightningModule):
    def __init__(self, model, loss_type, lr):
        super().__init__()
        self.model = model
        self.hparams['lr'] = lr
        self.loss_type = loss_type
        self.loss = nn.CrossEntropyLoss()
        self.save_hyperparameters()

    def forward(self, x):
        outputs = self.model(x)
        return outputs

    def training_step(self, batch, batch_idx):
        query_embeddings, target, all_batch_data, all_negative_pass = batch
        batch_size = target.shape[0]
        outputs = self.model(all_batch_data)
        target = torch.tensor(target, dtype=torch.float32, device=self.device)
        new_batch = []
        for ind, i in enumerate(outputs[0, :]):
            query_vector = i
            target_vector = target[ind]
            sim = torch.dot(query_vector, target_vector)
            all_negative_sim = torch.matmul(query_vector.unsqueeze(0), all_negative_pass.T)
            new_batch.append(torch.cat((sim.unsqueeze(0), all_negative_sim[0])).unsqueeze(0))
        new_batch = torch.cat(new_batch, dim=0)
        direct_target = torch.zeros(batch_size, device=self.device).long()
        contrastive_loss = self.loss(new_batch, direct_target)

        if self.loss_type == 'bidirectional':
            raise NotImplementedError()
        elif self.loss_type == 'direct':
            loss = contrastive_loss
            self.log('train/loss', loss, on_step=True, on_epoch=False)
            self.log('train/lr', self.hparams['lr'], on_step=True, on_epoch=False)
        elif self.loss_type == 'triplet':
            loss = contrastive_loss
        else:
            raise NotImplementedError()

        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=self.hparams['lr'])
        return optimizer

    def save(self, path):
        self.model.save_pretrained(path)