make serving tree use tree_output_dir as directory, node_edge_output_…

…file may be odps table

docs/source/quick_start/local_tutorial_tdm.md

@@ -43,15 +43,18 @@ python -m tzrec.tools.tdm.init_tree \
 --cate_id_field cate_id \
 --attr_fields cate_id,campaign_id,customer,brand,price \
 --node_edge_output_file data/init_tree
+--tree_output_dir data/init_tree
 ```
 
 - --item_input_path: 建树用的item特征文件
 - --item_id_field: 代表item的id的列名
 - --cate_id_field: 代表item的类别的列名
 - --attr_fields: (可选) 除了item_id外的item非数值型特征列名, 用逗号分开. 注意和配置文件中tdm_sampler顺序一致
 - --raw_attr_fields: (可选) item的数值型特征列名, 用逗号分开. 注意和配置文件中tdm_sampler顺序一致
-- --tree_output_file: (可选)初始树的保存路径, 不输入不会保存
-- --node_edge_output_file: 根据树生成的node和edge表的保存路径, 支持ODPS和本地txt两种
+- --node_edge_output_file: 根据树生成的node和edge表的保存路径, 支持`ODPS GL表`和`本地txt GL`两种
+  - ODPS GL表：设置形如`odps://{project}/tables/{tb_prefix}`，将会产出用于TDM训练负采样的GL Node表`odps://{project}/tables/{tb_prefix}_node_table`、GL Edge表`odps://{project}/tables/{tb_prefix}_edge_table`、用于离线检索的GL Edge表`odps://{project}/tables/{tb_prefix}_predict_edge_table`
+  - 本地txt GL表：设置的为目录， 将在目录下产出用于TDM训练负采样的GL Node表`node_table.txt`,GL Edge表`edge_table.txt`、用于离线检索的GL Edge表`predict_edge_table.txt`
+- --tree_output_dir: (可选) 树的保存目录, 将会在目录下存储`serving_tree`文件用于线上服务
 - --n_cluster: (可选,默认为2)树的分叉数
 
 #### 训练
@@ -115,6 +118,7 @@ OMP_NUM_THREADS=4 python tzrec/tools/tdm/cluster_tree.py \
     --embedding_field item_emb \
     --attr_fields cate_id,campaign_id,customer,brand,price \
     --node_edge_output_file data/learnt_tree \
+    --tree_output_dir data/learnt_tree \
     --parallel 16
 ```
 
@@ -123,8 +127,8 @@ OMP_NUM_THREADS=4 python tzrec/tools/tdm/cluster_tree.py \
 - --embedding_field: 代表item embedding的列名
 - --attr_fields: (可选) 除了item_id外的item非数值型特征列名, 用逗号分开. 注意和配置文件中tdm_sampler顺序一致
 - --raw_attr_fields: (可选) item的数值型特征列名, 用逗号分开. 注意和配置文件中tdm_sampler顺序一致
-- --tree_output_file: (可选)树的保存路径, 不输入不会保存
-- --node_edge_output_file: 根据树生成的node和edge表的保存路径, 支持ODPS和本地txt两种
+- --node_edge_output_file: 根据树生成的node和edge表的保存路径, 支持`ODPS GL表`和`本地txt GL`两种，同初始树
+- --tree_output_dir: (可选) 树的保存目录, 将会在目录下存储`serving_tree`文件用于线上服务
 - --n_cluster: (可选,默认为2)树的分叉数
 - --parllel: (可选，默认为16)聚类时CPU并行数
 

tzrec/tests/utils.py

@@ -792,6 +792,7 @@ def load_config_for_test(
                 f"--attr_fields {','.join(attr_fields)} "
                 f"--raw_attr_fields {','.join(raw_attr_fields)} "
                 f"--node_edge_output_file {test_dir}/init_tree "
+                f"--tree_output_dir {test_dir}/init_tree "
             )
             p = misc_util.run_cmd(cmd_str, os.path.join(test_dir, "log_init_tree.txt"))
             p.wait(600)
@@ -1212,6 +1213,7 @@ def test_tdm_cluster_train_eval(
         f"--attr_fields {','.join(attr_fields)} "
         f"--raw_attr_fields {','.join(raw_attr_fields)} "
         f"--node_edge_output_file {os.path.join(test_dir, 'learnt_tree')} "
+        f"--tree_output_dir {os.path.join(test_dir, 'learnt_tree')} "
         f"--parallel 1 "
     )
     p = misc_util.run_cmd(

tzrec/tools/tdm/cluster_tree.py

@@ -48,7 +48,7 @@
         help="The column names representing the raw features of item in the file.",
     )
     parser.add_argument(
-        "--tree_output_file",
+        "--tree_output_dir",
         type=str,
         default=None,
         help="The tree output file.",
@@ -78,16 +78,12 @@
         item_id_field=args.item_id_field,
         attr_fields=args.attr_fields,
         raw_attr_fields=args.raw_attr_fields,
-        output_file=args.tree_output_file,
+        output_dir=args.tree_output_dir,
         embedding_field=args.embedding_field,
         parallel=args.parallel,
         n_cluster=args.n_cluster,
     )
-    if args.tree_output_file:
-        save_tree = True
-    else:
-        save_tree = False
-    root = cluster.train(save_tree)
+    root = cluster.train()
     logger.info("Tree cluster done. Start save nodes and edges table.")
     tree_search = TreeSearch(
         output_file=args.node_edge_output_file,
@@ -96,5 +92,6 @@
     )
     tree_search.save()
     tree_search.save_predict_edge()
-    tree_search.save_serving_tree()
+    if args.tree_output_dir:
+        tree_search.save_serving_tree(args.tree_output_dir)
     logger.info("Save nodes and edges table done.")

tzrec/tools/tdm/gen_tree/tree_cluster.py

@@ -31,12 +31,12 @@ class TreeCluster:
     """Cluster based on emb vec.
 
     Args:
-        item_input_path(str): The file path where the item information is stored.
-        item_id_field(str): The column name representing item_id in the file.
-        attr_fields(List[str]): The column names representing the features in the file.
-        output_file(str): The output file.
-        parallel(int): The number of CPU cores for parallel processing.
-        n_cluster(int): The branching factor of the nodes in the tree.
+        item_input_path (str): The file path where the item information is stored.
+        item_id_field (str): The column name representing item_id in the file.
+        attr_fields (List[str]): The column names representing the features in the file.
+        output_dir (str): The output file.
+        parallel (int): The number of CPU cores for parallel processing.
+        n_cluster (int): The branching factor of the nodes in the tree.
     """
 
     def __init__(
@@ -45,7 +45,7 @@ def __init__(
         item_id_field: str,
         attr_fields: Optional[str] = None,
         raw_attr_fields: Optional[str] = None,
-        output_file: Optional[str] = None,
+        output_dir: Optional[str] = None,
         embedding_field: str = "item_emb",
         parallel: int = 16,
         n_cluster: int = 2,
@@ -60,7 +60,7 @@ def __init__(
         self.queue = None
         self.timeout = 5
         self.codes = None
-        self.output_file = output_file
+        self.output_dir = output_dir
         self.n_clusters = n_cluster
 
         self.item_id_field = item_id_field
@@ -140,7 +140,7 @@ def train(self, save_tree: bool = False) -> TDMTreeClass:
             p.join()
 
         assert queue.empty()
-        builder = tree_builder.TreeBuilder(self.output_file, self.n_clusters)
+        builder = tree_builder.TreeBuilder(self.output_dir, self.n_clusters)
         root = builder.build(
             self.ids, self.codes, self.attrs, self.raw_attrs, self.data, save_tree
         )

tzrec/tools/tdm/gen_tree/tree_generator.py

@@ -37,7 +37,7 @@ def __init__(
         cate_id_field: str,
         attr_fields: Optional[str] = None,
         raw_attr_fields: Optional[str] = None,
-        tree_output_file: Optional[str] = None,
+        tree_output_dir: Optional[str] = None,
         n_cluster: int = 2,
     ) -> None:
         self.item_input_path = item_input_path
@@ -49,7 +49,7 @@ def __init__(
             self.attr_fields = [x.strip() for x in attr_fields.split(",")]
         if raw_attr_fields:
             self.raw_attr_fields = [x.strip() for x in raw_attr_fields.split(",")]
-        self.tree_output_file = tree_output_file
+        self.tree_output_dir = tree_output_dir
         self.n_cluster = n_cluster
 
     def generate(self, save_tree: bool = False) -> TDMTreeClass:
@@ -141,6 +141,6 @@ def gen_code(start: int, end: int, code: int, items: List[Item]) -> None:
         )
         data = np.array([[] for i in range(len(ids))])
 
-        builder = TreeBuilder(self.tree_output_file, self.n_cluster)
+        builder = TreeBuilder(self.tree_output_dir, self.n_cluster)
         root = builder.build(ids, codes, attrs, raw_attrs, data, save_tree)
         return root

tzrec/tools/tdm/gen_tree/tree_search_util.py

@@ -195,9 +195,11 @@ def save_predict_edge(self) -> None:
                         for child in node.children:
                             f.write(f"{node.item_id}\t{child.item_id}\t{1.0}\n")
 
-    def save_serving_tree(self) -> None:
+    def save_serving_tree(self, tree_output_dir: str) -> None:
         """Save tree info for serving."""
-        with open(os.path.join(self.output_file, "serving_tree"), "w") as f:
+        if not os.path.exists(tree_output_dir):
+            os.makedirs(tree_output_dir)
+        with open(os.path.join(tree_output_dir, "serving_tree"), "w") as f:
             f.write(f"{self.max_level + 1} {self.child_num}\n")
             for _, nodes in enumerate(self.level_code):
                 for node in nodes:

tzrec/tools/tdm/gen_tree/tree_search_util_test.py

@@ -52,7 +52,7 @@ def test_cluster(self) -> None:
         search = TreeSearch(output_file=self.test_dir, root=root, child_num=2)
         search.save()
         search.save_predict_edge()
-        search.save_serving_tree()
+        search.save_serving_tree(self.test_dir)
 
         node_table = []
         edge_table = []

tzrec/tools/tdm/init_tree.py

@@ -48,10 +48,10 @@
         help="The column names representing the raw features of item in the file.",
     )
     parser.add_argument(
-        "--tree_output_file",
+        "--tree_output_dir",
         type=str,
         default=None,
-        help="The tree output file.",
+        help="The tree output directory.",
     )
     parser.add_argument(
         "--node_edge_output_file",
@@ -73,14 +73,10 @@
         cate_id_field=args.cate_id_field,
         attr_fields=args.attr_fields,
         raw_attr_fields=args.raw_attr_fields,
-        tree_output_file=args.tree_output_file,
+        tree_output_dir=args.tree_output_dir,
         n_cluster=args.n_cluster,
     )
-    if args.tree_output_file:
-        save_tree = True
-    else:
-        save_tree = False
-    root = generator.generate(save_tree)
+    root = generator.generate()
     logger.info("Tree init done. Start save nodes and edges table.")
     tree_search = TreeSearch(
         output_file=args.node_edge_output_file,
@@ -89,5 +85,6 @@
     )
     tree_search.save()
     tree_search.save_predict_edge()
-    tree_search.save_serving_tree()
+    if args.tree_output_dir:
+        tree_search.save_serving_tree(args.tree_output_dir)
     logger.info("Save nodes and edges table done.")