Skip to content

Instantly share code, notes, and snippets.

@jeanmidevacc

jeanmidevacc/surprise_get_closest_neighbours.py

Created March 15, 2022 12:14
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save jeanmidevacc/4d3ad98b27365b04e646d6862f4c2dc0 to your computer and use it in GitHub Desktop.
Save jeanmidevacc/4d3ad98b27365b04e646d6862f4c2dc0 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
surprise_get_closest_neighbours.py
def get_closest_neighbors(model, entityid, k, type_="item"):
if type_ == "item":
inner_entity_id = model.trainset.to_inner_iid(entityid)
else:
inner_entity_id = model.trainset.to_inner_uid(entityid)
closest_entity_id = model.get_neighbors(inner_entity_id, k)
if type_ == "item":
return [model.trainset.to_raw_iid(id_) for id_ in closest_entity_id]
else:
return [model.trainset.to_raw_uid(id_) for id_ in closest_entity_id]
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment