From 95f8c1577a3d7460c6f2b7a93bb65de8dfe826db Mon Sep 17 00:00:00 2001
From: zufchan <70619767+zufchan@users.noreply.github.com>
Date: Sat, 30 Jul 2022 18:28:18 +0600
Subject: [PATCH] Optimization for big values of num_list_per_user

I noticed that the current algorithm works slowly for big values of num_list_per_user, as a proposal, substitute a regular python list with a NumPy array for faster indexing. You can check the difference in https://colab.research.google.com/drive/1PMUomKqlEe48kzCeIZqQWNB0Rm_bLnJd?usp=sharing
---
 tensorflow_recommenders/examples/movielens.py | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/tensorflow_recommenders/examples/movielens.py b/tensorflow_recommenders/examples/movielens.py
index 02a0ba36..87ccc81a 100644
--- a/tensorflow_recommenders/examples/movielens.py
+++ b/tensorflow_recommenders/examples/movielens.py
@@ -112,13 +112,10 @@ def _sample_list(
       size=num_examples_per_list,
       replace=False,
   )
-  sampled_movie_titles = [
-      feature_lists["movie_title"][idx] for idx in sampled_indices
-  ]
-  sampled_ratings = [
-      feature_lists["user_rating"][idx]
-      for idx in sampled_indices
-  ]
+  
+  sampled_movie_titles = feature_lists["movie_title"][sampled_indices].tolist()
+  
+  sampled_ratings = feature_lists["user_rating"][sampled_indices].tolist()
 
   return (
       tf.concat(sampled_movie_titles, 0),
@@ -174,6 +171,8 @@ def sample_listwise(
   tensor_slices = {"user_id": [], "movie_title": [], "user_rating": []}
 
   for user_id, feature_lists in example_lists_by_user.items():
+    feature_lists["movie_title"] = np.array(feature_lists["movie_title"])
+    feature_lists["user_rating"] = np.array(feature_lists["user_rating"])
     for _ in range(num_list_per_user):
 
       # Drop the user if they don't have enough ratings.