Python中怎么利用Faiss库实现向量近邻搜索,相信很多没有经验的人对此束手无策,为此本文总结了问题出现的原因和解决方法,通过这篇文章希望你能解决这个问题。
1、读取预训练好的Embedding数据
import pandas as pd
import numpy as np
df = pd.read_csv("./datas/movielens_sparkals_item_embedding.csv")
# 提取要使用的电影ID列表,注意要转换成int64
ids = df["id"].values.astype(np.int64)
# 记录ID列表的大小
ids_size = ids.shape[0]
# 读取内容embedding数据,转换成二维array
import json
import numpy as np
datas = []
for x in df["features"]:
datas.append(json.loads(x))
# 变成二维array
datas = np.array(datas).astype(np.float32)
# 记录数据维度
dimension = datas.shape[1]
2、使用faiss建立索引
import faiss
index = faiss.IndexFlatL2(dimension)
index2 = faiss.IndexIDMap(index)
index2.add_with_ids(datas, ids)
3、实现近邻搜索
# 读取user embedding数据
df_user = pd.read_csv("./datas/user_embedding.csv")
df_user.head()
# 挑选一条user的embedding,转换成1行N列的二维array
user_embedding = np.array(json.loads(df_user[df_user["id"] == 10]["features"].iloc[0]))
user_embedding = np.expand_dims(user_embedding, axis=0).astype(np.float32)
user_embedding
# 实现搜索,这里的I就是近邻ID列表
tok = 30
D, I = index.search(user_embedding, topk) # actual search
4、拼接得到内容列表
# 把搜索出来的ID变成Series
target_ids = pd.Series(I[0], name="MovieID")
# 读取电影数据
df_movie = pd.read_csv("./datas/ml-1m/movies.dat",
sep="::", header=None, engine="python",
names = "MovieID::Title::Genres".split("::"))
# 实现内容JOIN
df_result = pd.merge(target_ids, df_movie)
看完上述内容,你们掌握Python中怎么利用Faiss库实现向量近邻搜索的方法了吗?如果还想学到更多技能或想了解更多相关内容,欢迎关注亿速云行业资讯频道,感谢各位的阅读!
免责声明:本站发布的内容(图片、视频和文字)以原创、转载和分享为主,文章观点不代表本网站立场,如果涉及侵权请联系站长邮箱:is@yisu.com进行举报,并提供相关证据,一经查实,将立刻删除涉嫌侵权内容。