A-A+

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

2023年05月12日 23:05 学习笔记 暂无评论 共942字 (阅读175 views次)

【注意:此文章为博主原创文章!转载需注意,请带原文链接,至少也要是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])

布施恩德可便相知重

微信扫一扫打赏

支付宝扫一扫打赏

×

给我留言