diff --git a/rectools/model_selection/cross_validate.py b/rectools/model_selection/cross_validate.py index e1e39c15..ebb986de 100644 --- a/rectools/model_selection/cross_validate.py +++ b/rectools/model_selection/cross_validate.py @@ -49,6 +49,7 @@ def cross_validate( # pylint: disable=too-many-locals k: int, filter_viewed: bool, items_to_recommend: tp.Optional[ExternalIds] = None, + users: tp.Optional[ExternalIds] = None, ) -> tp.Dict[str, tp.Any]: """ Run cross validation on multiple models with multiple metrics. @@ -115,6 +116,9 @@ def cross_validate( # pylint: disable=too-many-locals fold_dataset = _gen_2x_internal_ids_dataset(interactions_df_train, dataset.user_features, dataset.item_features) interactions_df_test = interactions.df.iloc[test_ids] # 1x internal + if users is not None: + internal_users = dataset.user_id_map.convert_to_internal(users, strict=False) + interactions_df_test = interactions_df_test[interactions_df_test[Columns.User].isin(internal_users)] test_users = interactions_df_test[Columns.User].unique() # 1x internal catalog = interactions_df_train[Columns.Item].unique() # 1x internal