def test_OnDiskDataset_TVTSet_ItemSet_node_pairs_labels_indexes():
    """Test TVTSet which returns ItemSet with node pairs and negative ones."""
    with tempfile.TemporaryDirectory() as test_dir:
        train_seeds = np.arange(2000).reshape(1000, 2)
        train_neg_dst = np.random.choice(1000 * 10, size=1000 * 10)
        train_neg_src = train_seeds[:, 0].repeat(10)
        train_neg_seeds = (
            np.concatenate((train_neg_dst, train_neg_src)).reshape(2, -1).T
        )
        train_seeds = np.concatenate((train_seeds, train_neg_seeds))
        train_seeds_path = os.path.join(test_dir, "train_seeds.npy")
        np.save(train_seeds_path, train_seeds)

        train_labels = torch.empty(1000 * 11)
        train_labels[:1000] = 1
        train_labels[1000:] = 0
        train_labels_path = os.path.join(test_dir, "train_labels.pt")
        torch.save(train_labels, train_labels_path)

        train_indexes = torch.arange(0, 1000)
        train_indexes = np.concatenate(
            (train_indexes, train_indexes.repeat_interleave(10))
        )
        train_indexes_path = os.path.join(test_dir, "train_indexes.pt")
        torch.save(train_indexes, train_indexes_path)

        validation_seeds = np.arange(2000, 4000).reshape(1000, 2)
        validation_neg_seeds = train_neg_seeds + 1
        validation_seeds = np.concatenate(
            (validation_seeds, validation_neg_seeds)
        )
        validation_seeds_path = os.path.join(test_dir, "validation_seeds.npy")
        np.save(validation_seeds_path, validation_seeds)
        validation_labels = train_labels
        validation_labels_path = os.path.join(test_dir, "validation_labels.pt")
        torch.save(validation_labels, validation_labels_path)

        validation_indexes = train_indexes
        validation_indexes_path = os.path.join(
            test_dir, "validation_indexes.pt"
        )
        torch.save(validation_indexes, validation_indexes_path)

        test_seeds = np.arange(4000, 6000).reshape(1000, 2)
        test_neg_seeds = train_neg_seeds + 2
        test_seeds = np.concatenate((test_seeds, test_neg_seeds))
        test_seeds_path = os.path.join(test_dir, "test_seeds.npy")
        np.save(test_seeds_path, test_seeds)
        test_labels = train_labels
        test_labels_path = os.path.join(test_dir, "test_labels.pt")
        torch.save(test_labels, test_labels_path)

        test_indexes = train_indexes
        test_indexes_path = os.path.join(test_dir, "test_indexes.pt")
        torch.save(test_indexes, test_indexes_path)

        yaml_content = f"""
            tasks:
              - name: link_prediction
                train_set:
                  - type: null
                    data:
                      - name: seeds
                        format: numpy
                        in_memory: true
                        path: {train_seeds_path}
                      - name: labels
                        format: torch
                        in_memory: true
                        path: {train_labels_path}
                      - name: indexes
                        format: torch
                        in_memory: true
                        path: {train_indexes_path}
                validation_set:
                  - data:
                      - name: seeds
                        format: numpy
                        in_memory: true
                        path: {validation_seeds_path}
                      - name: labels
                        format: torch
                        in_memory: true
                        path: {validation_labels_path}
                      - name: indexes
                        format: torch
                        in_memory: true
                        path: {validation_indexes_path}
                test_set:
                  - type: null
                    data:
                      - name: seeds
                        format: numpy
                        in_memory: true
                        path: {test_seeds_path}
                      - name: labels
                        format: torch
                        in_memory: true
                        path: {test_labels_path}
                      - name: indexes
                        format: torch
                        in_memory: true
                        path: {test_indexes_path}
        """
        dataset = write_yaml_and_load_dataset(yaml_content, test_dir)

        # Verify train set.
        train_set = dataset.tasks[0].train_set
        assert len(train_set) == 1000 * 11
        assert isinstance(train_set, gb.ItemSet)
        for i, (node_pair, label, index) in enumerate(train_set):
            assert node_pair[0] == train_seeds[i][0]
            assert node_pair[1] == train_seeds[i][1]
            assert label == train_labels[i]
            assert index == train_indexes[i]
        assert train_set.names == ("seeds", "labels", "indexes")
        train_set = None

        # Verify validation set.
        validation_set = dataset.tasks[0].validation_set
        assert len(validation_set) == 1000 * 11
        assert isinstance(validation_set, gb.ItemSet)
        for i, (node_pair, label, index) in enumerate(validation_set):
            assert node_pair[0] == validation_seeds[i][0]
            assert node_pair[1] == validation_seeds[i][1]
            assert label == validation_labels[i]
            assert index == validation_indexes[i]
        assert validation_set.names == ("seeds", "labels", "indexes")
        validation_set = None

        # Verify test set.
        test_set = dataset.tasks[0].test_set
        assert len(test_set) == 1000 * 11
        assert isinstance(test_set, gb.ItemSet)
        for i, (node_pair, label, index) in enumerate(test_set):
            assert node_pair[0] == test_seeds[i][0]
            assert label == test_labels[i]
            assert index == test_indexes[i]
        assert test_set.names == ("seeds", "labels", "indexes")
        test_set = None
        dataset = None
Test TVTSet which returns ItemSet with node pairs and negative ones.
tests/python/pytorch/graphbolt/impl/test_ondisk_dataset.py