item2vec by ids

word_vector/api.py

@@ -12,7 +12,7 @@ from utils.es import es_scan
 from utils.files import DATA_PATH, MODEL_PATH
 from utils.message import (send_msg_to_dingtalk, send_performance_msg_to_dingtalk)
 
-from word_vector.tractate import TRACTATE_CLICK_IDS_MODEL
+from word_vector.tractate import TRACTATE_CLICK_IDS, TRACTATE_CLICK_IDS_MODEL
 
 model_output_name = "w2v_model"
 model_path = os.path.join(MODEL_PATH, model_output_name)
@@ -97,7 +97,32 @@ def clicked_tractate_ids_item2vec_model(id, n=5):
         return res
     except KeyError as e:
         # send_msg_to_dingtalk("tractate_item2vec: " + str(e))
-        pass
+        return []
+    except Exception as e:
+        send_msg_to_dingtalk(str(traceback.format_exc()))
+        return []
+
+
+@bind("strategy_embedding/word_vector/tractate_item2vec_by_ids")
+def clicked_tractate_ids_item2vec_model_by_ids(ids, n=5):
+    """
+    ids: ["417522", "179986", "75804", "84682", "400022"]
+    return: [[('77419', 0.8857042193412781),
+              ('75761', 0.7772561311721802)],
+             [('84630', 0.9709808826446533),
+              ('84671', 0.9621062278747559)]]
+    """
+    try:
+        time_begin = time.time()
+        res = []
+        for id in ids:
+            if id in TRACTATE_CLICK_IDS:
+                res.extend(clicked_tractate_ids_item2vec_model(id, n))
+        time_end = time.time() - time_begin
+        if time_end > 0.04:
+            send_performance_msg_to_dingtalk("clicked_tractate_ids_item2vec_model_by_ids id={} n={} cost {:.3f}ms".format(
+                ids, n, time_end * 1000))
+        return res
     except Exception as e:
         send_msg_to_dingtalk(str(traceback.format_exc()))
         return []

word_vector/tractate.py

@@ -15,6 +15,7 @@ tractate_click_ids_model_path = os.path.join(MODEL_PATH, "tractate_click_ids_ite
 
 try:
     TRACTATE_CLICK_IDS_MODEL = word2vec.Word2Vec.load(tractate_click_ids_model_path)
+    TRACTATE_CLICK_IDS = set(TRACTATE_CLICK_IDS_MODEL.wv.vocab.keys())
 except Exception as e:
     print(e)