优化推荐系统和冷启动
Change-Id: I93d3091f249f2396a25702e01eb8dd5a9e95e8bc
diff --git a/recommend/utils/data_generator.py b/recommend/utils/graph_build.py
similarity index 92%
rename from recommend/utils/data_generator.py
rename to recommend/utils/graph_build.py
index d3ef9cf..becf5df 100644
--- a/recommend/utils/data_generator.py
+++ b/recommend/utils/graph_build.py
@@ -9,7 +9,7 @@
SqlPassword = "123456"
-def fetch_data():
+def fetch_user_seed_data():
conn = pymysql.connect(
host=SqlURL,
port=SqlPort,
@@ -74,8 +74,10 @@
f.write(f"{uid}\t{items}\t{times}\n")
-def build_user_seed_graph():
- download_rows, favorite_rows = fetch_data()
+def build_user_seed_graph(return_mapping=False):
+ download_rows, favorite_rows = fetch_user_seed_data()
records, user_set, seed_set = process_records(download_rows, favorite_rows)
user2idx, seed2idx = build_id_maps(user_set, seed_set)
group_and_write(records, user2idx, seed2idx)
+ if return_mapping:
+ return user2idx, seed2idx
\ No newline at end of file