diff --git a/implicit/nearest_neighbours.py b/implicit/nearest_neighbours.py index a1f36c1..29e50c2 100644 --- a/implicit/nearest_neighbours.py +++ b/implicit/nearest_neighbours.py @@ -62,6 +62,7 @@ def recommend( userid, user_items=user_items, N=N, + score_dtype=np.float64, filter_already_liked_items=filter_already_liked_items, filter_items=filter_items, recalculate_user=recalculate_user, @@ -115,7 +116,12 @@ def similar_items( if not np.isscalar(itemid): return _batch_call( - self.similar_items, itemid, N=N, filter_items=filter_items, items=items + self.similar_items, + itemid, + N=N, + score_dtype=np.float64, + filter_items=filter_items, + items=items, ) if filter_items is not None and items is not None: diff --git a/implicit/utils.py b/implicit/utils.py index f61da18..6b7ac5a 100644 --- a/implicit/utils.py +++ b/implicit/utils.py @@ -103,11 +103,11 @@ def augment_inner_product_matrix(factors): return max_norm, np.append(factors, extra_dimension.reshape(norms.shape[0], 1), axis=1) -def _batch_call(func, ids, *args, N=10, **kwargs): +def _batch_call(func, ids, *args, N=10, id_dtype=np.int32, score_dtype=np.float32, **kwargs): # we're running in batch mode, just loop over each item and call the scalar version of the # function - output_ids = np.zeros((len(ids), N), dtype=np.int32) - output_scores = np.zeros((len(ids), N), dtype=np.float32) + output_ids = np.zeros((len(ids), N), dtype=id_dtype) + output_scores = np.zeros((len(ids), N), dtype=score_dtype) user_items = kwargs.pop("user_items") if "user_items" in kwargs else None item_users = kwargs.pop("item_users") if "item_users" in kwargs else None