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}