import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Baseline: Cosine Similarity + Contrastive Loss
    """
    def __init__(self, margin=0.5):
        super(Model, self).__init__()
        self.margin = margin
    
    def forward(self, x: torch.Tensor, y: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        # 1. 计算余弦相似度
        norm_x = torch.sqrt(torch.sum(x * x, dim=1, keepdim=True))
        norm_y = torch.sqrt(torch.sum(y * y, dim=1, keepdim=True))
        dot_product = torch.sum(x * y, dim=1, keepdim=True)
        cosine_sim = dot_product / (norm_x * norm_y + 1e-8)
        cosine_sim = cosine_sim.squeeze(1)

        # 2. 计算对比损失
        loss_positive = labels * (1 - cosine_sim)
        loss_negative = (1 - labels) * torch.relu(cosine_sim - self.margin)
        loss = loss_positive + loss_negative
        
        return torch.sum(loss)

batch_size = 1024
feature_dim = 512

def get_inputs():
    x = torch.randn(batch_size, feature_dim)
    y = torch.randn(batch_size, feature_dim)
    # 生成0或1的标签
    labels = torch.randint(0, 2, (batch_size,)).float()
    return [x, y, labels]

def get_init_inputs():
    return []
