unike 3.0.1__py3-none-any.whl

unike/data/BernSampler.py

@@ -0,0 +1,168 @@
+# coding:utf-8
+#
+# unike/data/BernSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 30, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Feb 25, 2024
+#
+# 该脚本定义了 BernSampler 类.
+
+"""
+BernSampler - 平移模型和语义匹配模型的训练集数据采样器。
+"""
+
+import torch
+import typing
+import random
+import collections
+import numpy as np
+from .TradSampler import TradSampler
+from typing_extensions import override
+
+class BernSampler(TradSampler):
+
+    """
+    平移模型和语义匹配模型的训练集 Bern 数据采样器（伯努利分布），如果想获得更详细的信息请访问 :ref:`TransH <transh>`。
+    """
+
+    def __init__(
+        self,
+        in_path: str = "./",
+        ent_file: str = "entity2id.txt",
+        rel_file: str = "relation2id.txt",
+        train_file: str = "train2id.txt",
+        batch_size: int | None = None,
+        neg_ent: int = 1):
+
+        """创建 BernSampler 对象。
+
+        :param in_path: 数据集目录
+        :type in_path: str
+        :param ent_file: entity2id.txt
+        :type ent_file: str
+        :param rel_file: relation2id.txt
+        :type rel_file: str
+        :param train_file: train2id.txt
+        :type train_file: str
+        :param batch_size: batch size 在该采样器中不起作用，只是占位符。
+        :type batch_size: int | None
+        :param neg_ent: 对于每一个正三元组, 构建的负三元组的个数, 替换 entity
+        :type neg_ent: int
+        """
+
+        super().__init__(
+            in_path=in_path,
+            ent_file=ent_file,
+            rel_file=rel_file,
+            train_file=train_file,
+            batch_size = batch_size,
+            neg_ent = neg_ent
+        )
+        
+        self.tph, self.hpt = self.get_tph_hpt()
+
+    def get_tph_hpt(self) -> tuple[collections.defaultdict[float], collections.defaultdict[float]]:
+
+        """计算 tph 和 hpt。
+        
+        :returns: tph 和 hpt
+        :rtype: tuple[collections.defaultdict[float], collections.defaultdict[float]]
+        """
+
+        h_of_r = collections.defaultdict(set)
+        t_of_r = collections.defaultdict(set)
+        freq_rel = collections.defaultdict(float)
+        tph = collections.defaultdict(float)
+        hpt = collections.defaultdict(float)
+        for h, r, t in self.train_triples:
+            freq_rel[r] += 1.0
+            h_of_r[r].add(h)
+            t_of_r[r].add(t)
+        for r in h_of_r:
+            tph[r] = freq_rel[r] / len(h_of_r[r])
+            hpt[r] = freq_rel[r] / len(t_of_r[r])
+        return tph, hpt
+
+    @override
+    def sampling(
+        self,
+        pos_triples: list[tuple[int, int, int]]) -> dict[str, typing.Union[str, torch.Tensor]]:
+
+        """平移模型和语义匹配模型的训练集 bern 的数据采样函数（伯努利分布）。
+        
+        :param pos_triples: 知识图谱中的正确三元组
+        :type pos_triples: list[tuple[int, int, int]]
+        :returns: 平移模型和语义匹配模型的训练数据
+        :rtype: dict[str, typing.Union[str, torch.Tensor]]
+        """
+
+        batch_data = {}
+        neg_ent_sample = []
+
+        batch_data['mode'] = 'bern'
+        for h, r, t in pos_triples:
+            neg_ent = self.__normal_batch(h, r, t, self.neg_ent)
+            neg_ent_sample += neg_ent
+        
+        batch_data["positive_sample"] = torch.LongTensor(np.array(pos_triples))
+        batch_data["negative_sample"] = torch.LongTensor(np.array(neg_ent_sample))
+
+        return batch_data
+
+    def __normal_batch(
+        self,
+        h: int,
+        r: int,
+        t: int,
+        neg_size: int) -> list[tuple[int, int, int]]:
+
+        """Bern 负采样函数
+        
+        :param h: 头实体
+        :type h: int
+        :param r: 关系
+        :type r: int
+        :param t: 尾实体
+        :type t: int
+        :param neg_size: 负三元组个数
+        :type neg_size: int
+        :returns: 负三元组中的头实体列表
+        :rtype: list[tuple[int, int, int]]
+        """
+
+        neg_size_h = 0
+        neg_size_t = 0
+        prob = self.hpt[r] / (self.hpt[r] + self.tph[r])
+        for _ in range(neg_size):
+            if random.random() < prob:
+                neg_size_t += 1
+            else:
+                neg_size_h += 1
+
+        res = []
+
+        neg_list_h = []
+        neg_cur_size = 0
+        while neg_cur_size < neg_size_h:
+            neg_tmp_h = self.corrupt_head(t, r, num_max=(neg_size_h - neg_cur_size) * 2)
+            neg_list_h.append(neg_tmp_h)
+            neg_cur_size += len(neg_tmp_h)
+        if neg_list_h != []:
+            neg_list_h = np.concatenate(neg_list_h)
+        
+        for hh in neg_list_h[:neg_size_h]:
+            res.append((hh, r, t))
+        
+        neg_list_t = []
+        neg_cur_size = 0
+        while neg_cur_size < neg_size_t:
+            neg_tmp_t = self.corrupt_tail(h, r, num_max=(neg_size_t - neg_cur_size) * 2)
+            neg_list_t.append(neg_tmp_t)
+            neg_cur_size += len(neg_tmp_t)
+        if neg_list_t != []:
+            neg_list_t = np.concatenate(neg_list_t)
+        
+        for tt in neg_list_t[:neg_size_t]:
+            res.append((h, r, tt))
+
+        return res

unike/data/CompGCNSampler.py

@@ -0,0 +1,140 @@
+# coding:utf-8
+#
+# unike/data/CompGCNSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 19, 2023
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 19, 2023
+#
+# 该脚本定义了 CompGCNSampler 类.
+
+"""
+CompGCNSampler - CompGCN 的数据采样器。
+"""
+
+import dgl
+import torch
+import typing
+import numpy as np
+from .RGCNSampler import RGCNSampler
+from typing_extensions import override
+
+class CompGCNSampler(RGCNSampler):
+    
+    """``CompGCN`` :cite:`CompGCN` 的训练数据采样器。
+
+    例子::
+
+        from unike.data import RGCNSampler, CompGCNSampler
+        from torch.utils.data import DataLoader
+
+        #: 训练数据采样器
+        train_sampler: typing.Union[typing.Type[RGCNSampler], typing.Type[CompGCNSampler]] = train_sampler(
+            in_path=in_path,
+            ent_file=ent_file,
+            rel_file=rel_file,
+            train_file=train_file,
+            batch_size=batch_size,
+            neg_ent=neg_ent
+        )
+
+        #: 训练集三元组
+        data_train: list[tuple[int, int, int]] = train_sampler.get_train()
+
+        train_dataloader = DataLoader(
+            data_train,
+            shuffle=True,
+            batch_size=batch_size,
+            num_workers=num_workers,
+            pin_memory=True,
+            drop_last=True,
+            collate_fn=train_sampler.sampling,
+        )
+    """
+
+    def __init__(
+        self,
+        in_path: str = "./",
+        ent_file: str = "entity2id.txt",
+        rel_file: str = "relation2id.txt",
+        train_file: str = "train2id.txt",
+        batch_size: int | None = None,
+        neg_ent: int = 1):
+
+        """创建 CompGCNSampler 对象。
+
+        :param in_path: 数据集目录
+        :type in_path: str
+        :param ent_file: entity2id.txt
+        :type ent_file: str
+        :param rel_file: relation2id.txt
+        :type rel_file: str
+        :param train_file: train2id.txt
+        :type train_file: str
+        :param batch_size: batch size
+        :type batch_size: int | None
+        :param neg_ent: 对于 CompGCN 不起作用。
+        :type neg_ent: int
+        """
+
+        super().__init__(
+            in_path=in_path,
+            ent_file=ent_file,
+            rel_file=rel_file,
+            train_file=train_file,
+            batch_size=batch_size,
+            neg_ent=neg_ent
+        )
+        
+        super().get_hr_train()
+        
+        self.graph, self.relation, self.norm = \
+            self.build_graph(self.ent_tol, np.array(self.t_triples).transpose(), -0.5)
+
+    @override
+    def sampling(
+        self,
+        pos_hr_t: list[tuple[tuple[int, int], list[int]]]) -> dict[str, typing.Union[dgl.DGLGraph, torch.Tensor]]:
+        
+        """``CompGCN`` :cite:`CompGCN` 的采样函数。
+
+        :param pos_triples: 知识图谱中的正确三元组
+        :type pos_triples: list[tuple[tuple[int, int], list[int]]]
+        :returns: ``CompGCN`` :cite:`CompGCN` 的训练数据
+        :rtype: dict[str, typing.Union[dgl.DGLGraph, torch.Tensor]]
+        """
+
+        batch_data = {}
+        
+        self.label = torch.zeros(self.batch_size, self.ent_tol)
+        self.triples  = torch.LongTensor([hr for hr , _ in pos_hr_t])
+        for id, hr_sample in enumerate([t for _ ,t in pos_hr_t]):
+            self.label[id][hr_sample] = 1
+
+        batch_data['sample']   = self.triples
+        batch_data['label']    = self.label
+        batch_data['graph']    = self.graph
+        batch_data['relation'] = self.relation
+        batch_data['norm']     = self.norm
+
+        return batch_data
+
+    @override
+    def node_norm_to_edge_norm(
+        self,
+        graph: dgl.DGLGraph,
+        node_norm: torch.Tensor) -> torch.Tensor:
+
+        """根据源节点和目标节点的度计算每条边的归一化系数。
+
+        :param graph: 子图的节点数
+        :type graph: dgl.DGLGraph
+        :param node_norm: 节点的归一化系数
+        :type node_norm: torch.Tensor
+        :returns: 边的归一化系数
+        :rtype: torch.Tensor
+        """
+        
+        graph.ndata['norm'] = node_norm
+        graph.apply_edges(lambda edges: {'norm': edges.dst['norm'] * edges.src['norm']})
+        norm = graph.edata.pop('norm').squeeze()
+        return norm

unike/data/CompGCNTestSampler.py

@@ -0,0 +1,84 @@
+# coding:utf-8
+#
+# unike/data/CompGCNTestSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 19, 2023
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 19, 2023
+#
+# 该脚本定义了 CompGCNTestSampler 类.
+
+"""
+CompGCNTestSampler - CompGCN 的测试数据采样器。
+"""
+
+from .CompGCNSampler import CompGCNSampler
+from .RGCNTestSampler import RGCNTestSampler
+
+class CompGCNTestSampler(RGCNTestSampler):
+
+    """``CompGCN`` :cite:`CompGCN` 的测试数据采样器。
+
+    例子::
+
+        from unike.data import RGCNTestSampler, CompGCNTestSampler
+        from torch.utils.data import DataLoader
+
+        #: 测试数据采样器
+        test_sampler: typing.Union[typing.Type[RGCNTestSampler], typing.Type[CompGCNTestSampler]] = test_sampler(
+            sampler=train_sampler,
+            valid_file=valid_file,
+            test_file=test_file,
+        )
+    
+        #: 验证集三元组
+        data_val: list[tuple[int, int, int]] = test_sampler.get_valid()
+        #: 测试集三元组
+        data_test: list[tuple[int, int, int]] = test_sampler.get_test()
+
+        val_dataloader = DataLoader(
+            data_val,
+            shuffle=False,
+            batch_size=test_batch_size,
+            num_workers=num_workers,
+            pin_memory=True,
+            collate_fn=test_sampler.sampling,
+        )
+
+        test_dataloader = DataLoader(
+            data_test,
+            shuffle=False,
+            batch_size=test_batch_size,
+            num_workers=num_workers,
+            pin_memory=True,
+            collate_fn=test_sampler.sampling,
+        )
+    """
+
+    def __init__(
+        self,
+        sampler: CompGCNSampler,
+        valid_file: str = "valid2id.txt",
+        test_file: str = "test2id.txt",
+        type_constrain: bool = True):
+
+        """创建 CompGCNTestSampler 对象。
+
+        :param sampler: 训练数据采样器。
+        :type sampler: CompGCNSampler
+        :param valid_file: valid2id.txt
+        :type valid_file: str
+        :param test_file: test2id.txt
+        :type test_file: str
+        :param type_constrain: 是否报告 type_constrain.txt 限制的测试结果
+        :type type_constrain: bool
+        """
+
+        super().__init__(
+            sampler=sampler,
+            valid_file = valid_file,
+            test_file = test_file,
+            type_constrain = type_constrain
+        )
+        
+        #: 幂
+        self.power: float = -0.5

unike/data/KGEDataLoader.py

@@ -0,0 +1,315 @@
+# coding:utf-8
+#
+# unike/data/KGEDataLoader.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 16, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Apr 27, 2024
+#
+# 为 KGE 模型读取数据.
+
+"""
+KGEDataLoader - KGE 模型读取数据集类。
+"""
+
+import typing
+from .UniSampler import UniSampler
+from .BernSampler import BernSampler
+from .RGCNSampler import RGCNSampler
+from .CompGCNSampler import CompGCNSampler
+from .TestSampler import TestSampler
+from .TradTestSampler import TradTestSampler
+from torch.utils.data import DataLoader
+
+class KGEDataLoader:
+
+    """KGE 模型数据加载器。
+
+    例子::
+
+        from unike.data import KGEDataLoader, BernSampler, TradTestSampler
+
+        dataloader = KGEDataLoader(
+        	in_path = "../../benchmarks/FB15K/", 
+        	batch_size = 8192,
+        	neg_ent = 25,
+        	test = True,
+        	test_batch_size = 256,
+        	num_workers = 16,
+        	train_sampler = BernSampler,
+        	test_sampler = TradTestSampler
+        )
+    """
+    
+    def __init__(
+        self,
+        in_path: str = "./",
+        ent_file: str = "entity2id.txt",
+        rel_file: str = "relation2id.txt",
+        train_file: str = "train2id.txt",
+        valid_file: str = "valid2id.txt",
+        test_file: str = "test2id.txt",
+        batch_size: int | None = None,
+        neg_ent: int = 1,
+        test: bool = False,
+        test_batch_size: int | None = None,
+        type_constrain: bool = True,
+        num_workers: int | None = None,
+        train_sampler: typing.Union[typing.Type[UniSampler], typing.Type[BernSampler], typing.Type[RGCNSampler], typing.Type[CompGCNSampler]] = BernSampler,
+        test_sampler: typing.Type[TestSampler] = TradTestSampler):
+
+        """创建 KGEDataLoader 对象。
+
+        :param in_path: 数据集目录
+        :type in_path: str
+        :param ent_file: entity2id.txt
+        :type ent_file: str
+        :param rel_file: relation2id.txt
+        :type rel_file: str
+        :param train_file: train2id.txt
+        :type train_file: str
+        :param valid_file: valid2id.txt
+        :type valid_file: str
+        :param test_file: test2id.txt
+        :type test_file: str
+        :param batch_size: batch size
+        :type batch_size: int | None
+        :param neg_ent: 对于每一个正三元组, 构建的负三元组的个数, 替换 entity；对于 CompGCN 不起作用。
+        :type neg_ent: int
+        :param test: 是否读取验证集和测试集
+        :type test: bool
+        :param test_batch_size: test batch size
+        :type test_batch_size: int | None
+        :param type_constrain: 是否报告 type_constrain.txt 限制的测试结果
+        :type type_constrain: bool
+        :param num_workers: 加载数据的进程数
+        :type num_workers: int
+        :param train_sampler: 训练数据采样器
+        :type train_sampler: typing.Union[typing.Type[UniSampler], typing.Type[BernSampler], typing.Type[RGCNSampler], typing.Type[CompGCNSampler]]
+        :param test_sampler: 测试数据采样器
+        :type test_sampler: typing.Type[TestSampler]
+        """
+
+        #: 数据集目录
+        self.in_path: str = in_path
+        #: entity2id.txt
+        self.ent_file: str = ent_file
+        #: relation2id.txt
+        self.rel_file: str = rel_file
+        #: train2id.txt
+        self.train_file: str = train_file
+        #: valid2id.txt
+        self.valid_file: str = valid_file
+        #: test2id.txt
+        self.test_file: str = test_file
+        #: batch size
+        self.batch_size: int = batch_size
+        #: 对于每一个正三元组, 构建的负三元组的个数, 替换 entity；对于 CompGCN 不起作用。
+        self.neg_ent: int = neg_ent
+        #: 是否读取验证集和测试集
+        self.test: bool = test
+        #: test batch size
+        self.test_batch_size: int = test_batch_size
+        #: 是否报告 type_constrain.txt 限制的测试结果
+        self.type_constrain: bool = type_constrain
+        #: 加载数据的进程数
+        self.num_workers: int = num_workers
+
+        #: 训练数据采样器
+        self.train_sampler: typing.Union[UniSampler, BernSampler, RGCNSampler, CompGCNSampler] = train_sampler(
+            in_path=self.in_path,
+            ent_file=self.ent_file,
+            rel_file=self.rel_file,
+            train_file=self.train_file,
+            batch_size=self.batch_size,
+            neg_ent=self.neg_ent
+        )
+
+        #: 训练集三元组
+        self.data_train: list[tuple[int, int, int]] = self.train_sampler.get_train()
+
+        if self.test:
+            #: 测试数据采样器
+            self.test_sampler: TestSampler = test_sampler(
+                sampler=self.train_sampler,
+                valid_file=self.valid_file,
+                test_file=self.test_file,
+                type_constrain=type_constrain
+            )
+        
+            #: 验证集三元组
+            self.data_val: list[tuple[int, int, int]] = self.test_sampler.get_valid()
+            #: 测试集三元组
+            self.data_test: list[tuple[int, int, int]] = self.test_sampler.get_test()
+    
+    def get_ent_tol(self) -> int:
+
+        """返回实体个数。
+        
+        :returns: 实体个数
+        :rtype: int
+        """
+
+        return self.train_sampler.ent_tol
+
+    def get_rel_tol(self) -> int:
+
+        """返回关系个数。
+        
+        :returns: 关系个数
+        :rtype: int
+        """
+
+        return self.train_sampler.rel_tol
+
+    def train_dataloader(self) -> DataLoader:
+
+        """返回训练数据加载器。
+        
+        :returns: 训练数据加载器
+        :rtype: torch.utils.data.DataLoader
+        """
+
+        return DataLoader(
+            self.data_train,
+            shuffle=True,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            drop_last=True,
+            collate_fn=self.train_sampler.sampling,
+        )
+            
+    def val_dataloader(self) -> DataLoader:
+
+        """返回验证数据加载器。
+        
+        :returns: 验证数据加载器
+        :rtype: torch.utils.data.DataLoader
+        """
+
+        return DataLoader(
+            self.data_val,
+            shuffle=False,
+            batch_size=self.test_batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            collate_fn=self.test_sampler.sampling,
+        )
+
+    def test_dataloader(self) -> DataLoader:
+
+        """返回测试数据加载器。
+        
+        :returns: 测试数据加载器
+        :rtype: torch.utils.data.DataLoader"""
+        
+        return DataLoader(
+            self.data_test,
+            shuffle=False,
+            batch_size=self.test_batch_size,
+            num_workers=self.num_workers,
+            pin_memory=True,
+            collate_fn=self.test_sampler.sampling,
+        )
+
+def get_kge_data_loader_hpo_config() -> dict[str, dict[str, typing.Any]]:
+    
+    """返回 :py:class:`KGEDataLoader` 的默认超参数优化配置。
+    
+    默认配置为::
+    
+        parameters_dict = {
+            'dataloader': {
+                'value': 'KGEDataLoader'
+            },
+            'in_path': {
+                'value': './'
+            },
+            'ent_file': {
+                'value': 'entity2id.txt'
+            },
+            'rel_file': {
+                'value': 'relation2id.txt'
+            },
+            'train_file': {
+                'value': 'train2id.txt'
+            },
+            'valid_file': {
+                'value': 'valid2id.txt'
+            },
+            'test_file': {
+                'value': 'test2id.txt'
+            },
+            'batch_size': {
+                'values': [512, 1024, 2048, 4096]
+            },
+            'neg_ent': {
+                'values': [1, 4, 16, 64]
+            },
+            'test_batch_size': {
+                'value': 30
+            },
+            'type_constrain': {
+                'value': True
+            },
+            'num_workers': {
+                'value': 16
+            },
+            'train_sampler': {
+                'value': 'BernSampler'
+            },
+            'test_sampler': {
+                'value': 'TradTestSampler'
+            }
+        }
+    
+    :returns: :py:class:`KGEDataLoader` 的默认超参数优化配置
+    :rtype: dict[str, dict[str, typing.Any]]
+    """
+    
+    parameters_dict = {
+        'dataloader': {
+            'value': 'KGEDataLoader'
+        },
+        'in_path': {
+            'value': './'
+        },
+        'ent_file': {
+            'value': 'entity2id.txt'
+        },
+        'rel_file': {
+            'value': 'relation2id.txt'
+        },
+        'train_file': {
+            'value': 'train2id.txt'
+        },
+        'valid_file': {
+            'value': 'valid2id.txt'
+        },
+        'test_file': {
+            'value': 'test2id.txt'
+        },
+        'batch_size': {
+            'values': [512, 1024, 2048, 4096]
+        },
+        'neg_ent': {
+            'values': [1, 4, 16, 64]
+        },
+        'test_batch_size': {
+            'value': 30
+        },
+        'type_constrain': {
+            'value': True
+        },
+        'num_workers': {
+            'value': 16
+        },
+        'train_sampler': {
+            'value': 'BernSampler'
+        },
+        'test_sampler': {
+            'value': 'TradTestSampler'
+        }
+    }
+    
+    return parameters_dict