from sklearn.metrics.pairwise import cosine_similarity cos_sim = cosine_similarity(df_test['embed'].tolist(), df['cat_embeding'].tolist()) indexes = np.argmax(cos_sim, axis=1) cats = df.loc[indexes]['cat_name'] df_test['prediction'] = cats.tolist()