172 lines
6.0 KiB
Python
172 lines
6.0 KiB
Python
# # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
|
|
# #
|
|
# # Licensed under the Apache License, Version 2.0 (the "License");
|
|
# # you may not use this file except in compliance with the License.
|
|
# # You may obtain a copy of the License at
|
|
# #
|
|
# # http://www.apache.org/licenses/LICENSE-2.0
|
|
# #
|
|
# # Unless required by applicable law or agreed to in writing, software
|
|
# # distributed under the License is distributed on an "AS IS" BASIS,
|
|
# # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# # See the License for the specific language governing permissions and
|
|
# # limitations under the License.
|
|
|
|
# import os
|
|
|
|
# os.environ['FLAGS_enable_pir_api'] = '0'
|
|
|
|
# import tempfile
|
|
# import unittest
|
|
|
|
# import paddle
|
|
# from paddle.dataset.common import download
|
|
# from paddle.distributed.fleet.dataset import TreeIndex
|
|
|
|
# paddle.enable_static()
|
|
|
|
|
|
# def create_feeds():
|
|
# user_input = paddle.static.data(
|
|
# name="item_id", shape=[-1, 1], dtype="int64"
|
|
# )
|
|
|
|
# item = paddle.static.data(name="unit_id", shape=[-1, 1], dtype="int64")
|
|
|
|
# label = paddle.static.data(name="label", shape=[-1, 1], dtype="int64")
|
|
# labels = paddle.static.data(name="labels", shape=[-1, 1], dtype="int64")
|
|
|
|
# feed_list = [user_input, item, label, labels]
|
|
# return feed_list
|
|
|
|
|
|
# class TestTreeIndex(unittest.TestCase):
|
|
# def test_tree_index(self):
|
|
# path = download(
|
|
# "https://paddlerec.bj.bcebos.com/tree-based/data/mini_tree.pb",
|
|
# "tree_index_unittest",
|
|
# "e2ba4561c2e9432b532df40546390efa",
|
|
# )
|
|
# '''
|
|
# path = download(
|
|
# "https://paddlerec.bj.bcebos.com/tree-based/data/mini_tree.pb",
|
|
# "tree_index_unittest", "cadec20089f5a8a44d320e117d9f9f1a")
|
|
# '''
|
|
# tree = TreeIndex("demo", path)
|
|
# height = tree.height()
|
|
# branch = tree.branch()
|
|
# self.assertTrue(height == 5)
|
|
# self.assertTrue(branch == 2)
|
|
# self.assertEqual(tree.total_node_nums(), 25)
|
|
# self.assertEqual(tree.emb_size(), 30)
|
|
|
|
# # get_layer_codes
|
|
# layer_node_ids = []
|
|
# layer_node_codes = []
|
|
# for i in range(tree.height()):
|
|
# layer_node_codes.append(tree.get_layer_codes(i))
|
|
# layer_node_ids.append(
|
|
# [node.id() for node in tree.get_nodes(layer_node_codes[-1])]
|
|
# )
|
|
|
|
# all_leaf_ids = [node.id() for node in tree.get_all_leaves()]
|
|
# self.assertEqual(sum(all_leaf_ids), sum(layer_node_ids[-1]))
|
|
|
|
# # get_travel
|
|
# travel_codes = tree.get_travel_codes(all_leaf_ids[0])
|
|
# travel_ids = [node.id() for node in tree.get_nodes(travel_codes)]
|
|
|
|
# for i in range(height):
|
|
# self.assertIn(travel_ids[i], layer_node_ids[height - 1 - i])
|
|
# self.assertIn(travel_codes[i], layer_node_codes[height - 1 - i])
|
|
|
|
# # get_ancestor
|
|
# ancestor_codes = tree.get_ancestor_codes([all_leaf_ids[0]], height - 2)
|
|
# ancestor_ids = [node.id() for node in tree.get_nodes(ancestor_codes)]
|
|
|
|
# self.assertEqual(ancestor_ids[0], travel_ids[1])
|
|
# self.assertEqual(ancestor_codes[0], travel_codes[1])
|
|
|
|
# # get_pi_relation
|
|
# pi_relation = tree.get_pi_relation([all_leaf_ids[0]], height - 2)
|
|
# self.assertEqual(pi_relation[all_leaf_ids[0]], ancestor_codes[0])
|
|
|
|
# # get_travel_path
|
|
# travel_path_codes = tree.get_travel_path(
|
|
# travel_codes[0], travel_codes[-1]
|
|
# )
|
|
# travel_path_ids = [
|
|
# node.id() for node in tree.get_nodes(travel_path_codes)
|
|
# ]
|
|
|
|
# self.assertEqual([*travel_path_ids, travel_ids[-1]], travel_ids)
|
|
# self.assertEqual([*travel_path_codes, travel_codes[-1]], travel_codes)
|
|
|
|
# # get_children
|
|
# children_codes = tree.get_children_codes(travel_codes[1], height - 1)
|
|
# children_ids = [node.id() for node in tree.get_nodes(children_codes)]
|
|
# self.assertIn(all_leaf_ids[0], children_ids)
|
|
|
|
|
|
# class TestIndexSampler(unittest.TestCase):
|
|
# def setUp(self):
|
|
# self.temp_dir = tempfile.TemporaryDirectory()
|
|
|
|
# def tearDown(self):
|
|
# self.temp_dir.cleanup()
|
|
|
|
# def test_layerwise_sampler(self):
|
|
# path = download(
|
|
# "https://paddlerec.bj.bcebos.com/tree-based/data/mini_tree.pb",
|
|
# "tree_index_unittest",
|
|
# "e2ba4561c2e9432b532df40546390efa",
|
|
# )
|
|
|
|
# tdm_layer_counts = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
|
|
# # tree = TreeIndex("demo", path)
|
|
# file_name = os.path.join(
|
|
# self.temp_dir.name, "test_in_memory_dataset_tdm_sample_run.txt"
|
|
# )
|
|
# with open(file_name, "w") as f:
|
|
# # data = "29 d 29 d 29 29 29 29 29 29 29 29 29 29 29 29\n"
|
|
# data = "1 1 1 15 15 15\n"
|
|
# data += "1 1 1 15 15 15\n"
|
|
# f.write(data)
|
|
|
|
# slots = ["slot1", "slot2", "slot3"]
|
|
# slots_vars = []
|
|
# for slot in slots:
|
|
# var = paddle.static.data(name=slot, shape=[-1, 1], dtype="int64")
|
|
# slots_vars.append(var)
|
|
|
|
# dataset = paddle.distributed.InMemoryDataset()
|
|
# dataset.init(
|
|
# batch_size=1,
|
|
# pipe_command="cat",
|
|
# download_cmd="cat",
|
|
# use_var=slots_vars,
|
|
# )
|
|
# dataset.set_filelist([file_name])
|
|
# # dataset.update_settings(pipe_command="cat")
|
|
# # dataset._init_distributed_settings(
|
|
# # parse_ins_id=True,
|
|
# # parse_content=True,
|
|
# # fea_eval=True,
|
|
# # candidate_size=10000)
|
|
|
|
# dataset.load_into_memory()
|
|
# dataset.tdm_sample(
|
|
# 'demo',
|
|
# tree_path=path,
|
|
# tdm_layer_counts=tdm_layer_counts,
|
|
# start_sample_layer=1,
|
|
# with_hierarchy=False,
|
|
# seed=0,
|
|
# id_slot=2,
|
|
# )
|
|
# self.assertTrue(dataset.get_shuffle_data_size() == 8)
|
|
|
|
|
|
# if __name__ == '__main__':
|
|
# unittest.main()
|