A-A+
使用 FAISS 计算字符串之间的相似度

【注意:此文章为博主原创文章!转载需注意,请带原文链接,至少也要是txt格式!】
代码如下:
import faiss
import numpy as np
from transformers import AutoTokenizer, AutoModel
# 加载text2vec-large-chinese模型和分词器
model_name = "bert-base-chinese"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)
# 假设您有两个字符串s1和s2
s1 = "北京是中国的首都"
s2 = "上海是中国的金融中心"
# 使用分词器将字符串转换为输入张量
inputs1 = tokenizer(s1, return_tensors="pt", padding=True, truncation=True)
inputs2 = tokenizer(s2, return_tensors="pt", padding=True, truncation=True)
# 使用text2vec-large-chinese模型将输入张量转换为向量表示
with torch.no_grad():
outputs1 = model(**inputs1)
outputs2 = model(**inputs2)
embeddings1 = outputs1.last_hidden_state.mean(dim=1).numpy()
embeddings2 = outputs2.last_hidden_state.mean(dim=1).numpy()
# 创建一个FAISS索引对象
d = embeddings1.shape[1] # 向量维度
index = faiss.IndexFlatL2(d) # 创建L2距离度量的Flat索引
# 将向量添加到索引中
xb = np.concatenate((embeddings1, embeddings2)).astype('float32')
index.add(xb)
# 计算两个向量之间的相似度
xq = np.array([xb[0]]) # 查询向量
k = 1 # 返回最相似的k个向量
D, I = index.search(xq, k)
# 输出相似度的度量值
print("相似度: ", 1 - D[0][0])
布施恩德可便相知重
微信扫一扫打赏
支付宝扫一扫打赏