Python中怎么利用Faiss庫實(shí)現(xiàn)向量近鄰搜索

import pandas as pdimport numpy as np
df = pd.read_csv("./datas/movielens_sparkals_item_embedding.csv")
# 提取要使用的電影ID列表，注意要轉(zhuǎn)換成int64ids = df["id"].values.astype(np.int64)
# 記錄ID列表的大小ids_size = ids.shape[0]
# 讀取內(nèi)容embedding數(shù)據(jù)，轉(zhuǎn)換成二維arrayimport jsonimport numpy as npdatas = []for x in df["features"]:    datas.append(json.loads(x))# 變成二維arraydatas = np.array(datas).astype(np.float32)
# 記錄數(shù)據(jù)維度dimension = datas.shape[1]

import faiss
index = faiss.IndexFlatL2(dimension)index2 = faiss.IndexIDMap(index)
index2.add_with_ids(datas, ids)

# 讀取user embedding數(shù)據(jù)df_user = pd.read_csv("./datas/user_embedding.csv")df_user.head()
# 挑選一條user的embedding，轉(zhuǎn)換成1行N列的二維arrayuser_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
# 實(shí)現(xiàn)搜索，這里的I就是近鄰ID列表tok = 30D, I = index.search(user_embedding, topk)     # actual search

# 把搜索出來的ID變成Seriestarget_ids = pd.Series(I[0], name="MovieID")
# 讀取電影數(shù)據(jù)df_movie = pd.read_csv("./datas/ml-1m/movies.dat",                     sep="::", header=None, engine="python",                     names = "MovieID::Title::Genres".split("::"))# 實(shí)現(xiàn)內(nèi)容JOINdf_result = pd.merge(target_ids, df_movie)