numpy 实现embedding 使用余弦相似度的NTxent

下面是一个使用numpy实现embedding使用余弦相似度的NTxent的示例代码：

import numpy as np

# 定义一些超参数
k = 10 # top-k正样本
T = 0.5 # 温度参数

# 假设我们有n个样本，每个样本有m维特征
n, m = 1000, 128

# 随机初始化嵌入向量
embeddings = np.random.randn(n, m)

# 计算余弦相似度矩阵
norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
similarities = np.dot(embeddings, embeddings.T) / np.dot(norms, norms.T)

# 对于每个样本i，选择top-k个正样本和一组负样本
pos_samples = np.zeros((n, k))
neg_samples = np.zeros((n, k))
for i in range(n):
    pos_samples[i] = np.argsort(similarities[i])[::-1][1:k+1]
    neg_samples[i] = np.random.choice(np.delete(np.arange(n), pos_samples[i]), k, replace=False)

# 计算交叉熵损失
loss = 0
for i in range(n):
    pos_scores = np.sum(embeddings[i] * embeddings[pos_samples[i]], axis=1) / T
    neg_scores = np.sum(embeddings[i] * embeddings[neg_samples[i]], axis=1) / T
    scores = np.concatenate((pos_scores, neg_scores))
    labels = np.zeros(k*2)
    labels[:k] = 1
    loss += -np.mean(labels * np.log(np.exp(scores) / np.sum(np.exp(scores))))
print('NTxent loss:', loss)

在这个示例中，我们首先随机初始化n个样本的m维嵌入向量，然后计算余弦相似度矩阵。接着，对于每个样本i，我们选择top-k个余弦相似度最大的样本作为正样本，随机选择一组与i余弦相似度较小的样本作为负样本。最后，我们计算每个样本i的交叉熵损失，并输出总的NTxent损失。

需要注意的是，这只是一个简单的示例，实际使用时需要根据具体情况进行调整和优化。此外，还需要对嵌入向量进行归一化处理，以保证余弦相似度的取值范围在[-1, 1]之间。