"""Model architecture for similarity prediction."""
import torch
import torch.nn as nn
import torch.nn.functional as F
from sentence_transformers import SentenceTransformer


class SimilarityModel(nn.Module):
    """Similarity model based on CodeRankEmbed."""
    
    def __init__(self, model_name="nomic-ai/CodeRankEmbed", dropout=0.1):
        super().__init__()
        st = SentenceTransformer(model_name)
        self.encoder = st._first_module().auto_model
        self.hidden_size = self.encoder.config.hidden_size
        self.projection = nn.Linear(self.hidden_size, self.hidden_size)
        self.dropout = nn.Dropout(dropout)
    
    def encode(self, input_ids, attention_mask):
        out = self.encoder(input_ids, attention_mask)
        emb = out[0]
        mask_exp = attention_mask.unsqueeze(-1).expand(emb.size()).float()
        pooled = torch.sum(emb * mask_exp, 1) / torch.clamp(mask_exp.sum(1), min=1e-9)
        pooled = self.dropout(pooled)
        proj = self.projection(pooled)
        return F.normalize(proj, p=2, dim=1)
    
    def forward(self, input_ids_1, attention_mask_1, input_ids_2, attention_mask_2):
        emb1 = self.encode(input_ids_1, attention_mask_1)
        emb2 = self.encode(input_ids_2, attention_mask_2)
        sim = torch.sum(emb1 * emb2, dim=1)
        return (sim + 1.0) / 2.0