add timer

src/main.py

 import os
+import random
 import shutil
 import time
 import timeit
@@ -17,6 +18,20 @@ from models.esmm.model import esmm_model_fn, model_export, model_predict
 # tf.compat.v1.enable_eager_execution()
 
 
+def time_cost(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        time_1 = timeit.default_timer()
+
+        ret = func(*args, **kwargs)
+
+        total_1 = (timeit.default_timer() - time_1)
+        print("prediction total cost {:.5f}s".format(total_1))
+        return ret
+
+    return wrapper
+
+
 def main():
     time_begin = time.time()
 
@@ -66,17 +81,25 @@ def main():
     df2 = get_diary_df_from_redis()
     redis_device_df = device_feature_engineering(df)
     redis_diary_df = diary_feature_engineering(df2, from_redis=True)
-    diary_ids = list(redis_diary_df["card_id"].values)[:300]
+    diary_ids = list(redis_diary_df["card_id"].values)
+
+    @time_cost
+    def test1():
+        user1 = join_device_diary(device_id, random.sample(diary_ids, 300), redis_device_df, redis_diary_df)
+        model_predict(user1, predict_fn)
+
+    test1()
 
-    time_1 = timeit.default_timer()
-    res = join_device_diary(device_id, diary_ids, redis_device_df, redis_diary_df)
-    print(len(res))
-    print(res.sample(1), "\n")
-    print(res.sample(1))
+    # time_1 = timeit.default_timer()
+    # user1 = join_device_diary(device_id, random.sample(diary_ids, 300), redis_device_df, redis_diary_df)
+    # model_predict(user1, predict_fn)
+    # total_1 = (timeit.default_timer() - time_1)
+    # print("prediction total cost {:.5f}s".format(total_1))
 
-    model_predict(res, predict_fn)
-    total_1 = (timeit.default_timer() - time_1)
-    print("prediction total cost {:.5f}s".format(total_1))
+    # user2 = join_device_diary(device_id, random.sample(diary_ids, 300), redis_device_df, redis_diary_df)
+    # user3 = join_device_diary(device_id, random.sample(diary_ids, 300), redis_device_df, redis_diary_df)
+    # user4 = join_device_diary(device_id, random.sample(diary_ids, 300), redis_device_df, redis_diary_df)
+    # user5 = join_device_diary(device_id, random.sample(diary_ids, 300), redis_device_df, redis_diary_df)
 
     total_time = (time.time() - time_begin) / 60
     print("cost {:.2f} mins at {}".format(total_time, datetime.now()))