|
|
from __future__ import annotations |
|
|
from transformers import BertModel, BertConfig, PreTrainedModel |
|
|
from transformers.tokenization_utils_base import BatchEncoding |
|
|
import torch, torch.nn as nn, torch.nn.functional as F |
|
|
|
|
|
class SimCSEInferenceModel(PreTrainedModel): |
|
|
config_class = BertConfig |
|
|
def __init__(self, config): |
|
|
super().__init__(config) |
|
|
base_cfg = BertConfig(**config.to_dict()) |
|
|
self.encoder_input = BertModel(base_cfg) |
|
|
self.encoder_output = BertModel(base_cfg) |
|
|
hidden = self.encoder_input.config.hidden_size |
|
|
self.dense_input = nn.Linear(hidden, hidden) |
|
|
self.dense_output = nn.Linear(hidden, hidden) |
|
|
self.activation = nn.Tanh() |
|
|
self.temperature = getattr(config, "simcse_temperature", 0.05) |
|
|
@torch.no_grad() |
|
|
def encode_input(self, tok: BatchEncoding): |
|
|
h = self.encoder_input(**tok).last_hidden_state[:, 0] |
|
|
return self.activation(self.dense_input(h)) |
|
|
@torch.no_grad() |
|
|
def encode_output(self, tok: BatchEncoding): |
|
|
h = self.encoder_output(**tok).last_hidden_state[:, 0] |
|
|
return self.activation(self.dense_output(h)) |
|
|
def forward(self, tokenized_texts_1, tokenized_texts_2, labels, **_): |
|
|
device = next(self.parameters()).device |
|
|
z1 = F.normalize(self.encode_input(tokenized_texts_1.to(device)), dim=-1) |
|
|
z2 = F.normalize(self.encode_output(tokenized_texts_2.to(device)), dim=-1) |
|
|
sim = torch.matmul(z1, z2.T) |
|
|
loss = F.cross_entropy(sim / self.temperature, labels.to(device)) |
|
|
return {"loss": loss, "logits": sim} |
|
|
|
