unike 3.0.1__py3-none-any.whl

unike/data/KGReader.py

@@ -0,0 +1,138 @@
+# coding:utf-8
+#
+# unike/data/KGReader.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 17, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Feb 25, 2024
+#
+# 从文件中读取知识图谱.
+
+"""
+KGReader - 从文件中读取知识图谱。
+"""
+
+import os
+import numpy as np
+import collections
+
+class KGReader:
+
+    """
+    从文件中读取知识图谱。
+    """
+    
+    def __init__(
+        self,
+        in_path: str = "./",
+        ent_file: str = "entity2id.txt",
+        rel_file: str = "relation2id.txt",
+        train_file: str = "train2id.txt"):
+
+        """创建 KGReader 对象。
+
+        :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
+        """
+
+        #: 数据集目录
+        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
+
+        #: 实体的个数
+        self.ent_tol: int = 0
+        #: 关系的个数
+        self.rel_tol: int = 0
+        #: 训练集三元组的个数
+        self.train_tol: int = 0
+
+        #: 实体->ID
+        self.ent2id: dict = {}
+        #: 关系->ID
+        self.rel2id: dict = {}
+        #: ID->实体
+        self.id2ent: dict = {}
+        #: ID->关系
+        self.id2rel: dict = {}
+
+        #: 训练集三元组
+        self.train_triples: list[tuple[int, int, int]] = []
+
+        #: 训练集中所有 h-r 对对应的 t 集合
+        self.hr2t_train: collections.defaultdict[set] = collections.defaultdict(set)
+        #: 训练集中所有 r-t 对对应的 h 集合
+        self.rt2h_train: collections.defaultdict[set] = collections.defaultdict(set)
+
+        self.get_id()
+        self.get_train_triples_id()
+    
+    def get_id(self):
+
+        """读取 :py:attr:`ent_file` 文件和 :py:attr:`rel_file` 文件。"""
+        
+        with open(os.path.join(self.in_path, self.ent_file)) as f:
+            self.ent_tol = (int)(f.readline())
+            for line in f:
+                entity, eid = line.strip().split("\t")
+                self.ent2id[entity] = int(eid)
+                self.id2ent[int(eid)] = entity
+        
+        with open(os.path.join(self.in_path, self.rel_file)) as f:
+            self.rel_tol = (int)(f.readline())
+            for line in f:
+                relation, rid = line.strip().split("\t")
+                self.rel2id[relation] = int(rid)
+                self.id2rel[int(rid)] = relation
+    
+    def get_train_triples_id(self):
+
+        """读取 :py:attr:`train_file` 文件。"""
+        
+        with open(os.path.join(self.in_path, self.train_file)) as f:
+            self.train_tol = (int)(f.readline())
+            for line in f:
+                h, t, r = line.strip().split()
+                self.train_triples.append((int(h), int(r), int(t)))
+        
+    def get_hr2t_rt2h_from_train(self):
+
+        """获得 :py:attr:`hr2t_train` 和 :py:attr:`rt2h_train` 。"""
+        
+        for h, r, t in self.train_triples:
+            self.hr2t_train[(h, r)].add(t)
+            self.rt2h_train[(r, t)].add(h)
+        for h, r in self.hr2t_train:
+            self.hr2t_train[(h, r)] = np.array(list(self.hr2t_train[(h, r)]))
+        for r, t in self.rt2h_train:
+            self.rt2h_train[(r, t)] = np.array(list(self.rt2h_train[(r, t)]))
+            
+    def get_hr_train(self):
+
+        """用于 ``CompGCN`` :cite:`CompGCN` 训练，因为 ``CompGCN`` :cite:`CompGCN` 的组合运算仅需要头实体和关系。
+        
+        如果想获得更详细的信息请访问 :ref:`CompGCN <compgcn>`。
+        """
+        
+        self.t_triples = self.train_triples
+        self.train_triples = [(hr, list(t)) for (hr,t) in self.hr2t_train.items()]
+        
+    def get_train(self) -> list[tuple[int, int, int]]:
+
+        """
+        返回训练集三元组。
+
+        :returns: :py:attr:`train_triples`
+        :rtype: list[tuple[int, int, int]]
+        """
+
+        return self.train_triples

unike/data/RGCNSampler.py

@@ -0,0 +1,261 @@
+# coding:utf-8
+#
+# unike/data/RGCNSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 16, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 21, 2024
+#
+# R-GCN 的数据采样器.
+
+"""
+RGCNSampler - R-GCN 的数据采样器。
+"""
+
+import dgl
+import torch
+import typing
+import warnings
+import numpy as np
+from .RevSampler import RevSampler
+
+warnings.filterwarnings("ignore")
+
+class RGCNSampler(RevSampler):
+
+    """``R-GCN`` :cite:`R-GCN` 的训练数据采样器。
+
+    例子::
+
+        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):
+
+        """创建 RGCNSampler 对象。
+
+        :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 (head + tail)
+        :type neg_ent: int
+        """
+
+        super().__init__(
+            in_path=in_path,
+            ent_file=ent_file,
+            rel_file=rel_file,
+            train_file=train_file
+        )
+
+        #: batch size
+        self.batch_size: int = batch_size
+        #: 对于每一个正三元组, 构建的负三元组的个数, 替换 entity (head + tail)
+        self.neg_ent: int = neg_ent
+
+        self.entity   = None
+        self.triples  = None
+        self.label    = None
+        self.graph    = None
+        self.relation = None
+        self.norm     = None
+
+    def sampling(
+        self,
+        pos_triples: list[tuple[int, int, int]]) -> dict[str, typing.Union[dgl.DGLGraph, torch.Tensor]]:
+
+        """``R-GCN`` :cite:`R-GCN` 的采样函数。
+        
+        :param pos_triples: 知识图谱中的正确三元组
+        :type pos_triples: list[tuple[int, int, int]]
+        :returns: ``R-GCN`` :cite:`R-GCN` 的训练数据
+        :rtype: dict[str, typing.Union[dgl.DGLGraph , torch.Tensor]]
+        """
+        
+        batch_data = {}
+        
+        pos_triples = np.array(pos_triples)
+        pos_triples, self.entity = self.sampling_positive(pos_triples)
+        head_triples = self.sampling_negative('head', pos_triples)
+        tail_triples = self.sampling_negative('tail', pos_triples)
+        self.triples = np.concatenate((pos_triples,head_triples,tail_triples))
+        batch_data['entity']  = self.entity
+        batch_data['triples'] = torch.from_numpy(self.triples)
+        
+        self.label = torch.zeros((len(self.triples),1))
+        self.label[0 : self.batch_size] = 1
+        batch_data['label'] = self.label
+        
+        split_size = int(self.batch_size * 0.5) 
+        graph_split_ids = np.random.choice(
+            self.batch_size,
+            size=split_size, 
+            replace=False
+        )
+        head,rela,tail = pos_triples.transpose()
+        head = torch.tensor(head[graph_split_ids], dtype=torch.long).contiguous()
+        rela = torch.tensor(rela[graph_split_ids], dtype=torch.long).contiguous()
+        tail = torch.tensor(tail[graph_split_ids], dtype=torch.long).contiguous()
+        self.graph, self.relation, self.norm = self.build_graph(len(self.entity), (head,rela,tail), -1)
+        batch_data['graph']    = self.graph
+        batch_data['relation'] = self.relation
+        batch_data['norm']     = self.norm
+
+        return batch_data
+
+    def sampling_positive(
+        self,
+        positive_triples: list[tuple[int, int, int]]) -> tuple[np.ndarray, torch.Tensor]:
+
+        """为创建子图重新采样三元组子集，重排实体 ID。
+        
+        :param pos_triples: 知识图谱中的正确三元组
+        :type pos_triples: list[tuple[int, int, int]]
+        :returns: 三元组子集和原始的实体 ID
+        :rtype: tuple[numpy.ndarray, torch.Tensor]
+        """
+
+        edges = np.random.choice(
+            np.arange(len(positive_triples)),
+            size = self.batch_size,
+            replace=False
+        )
+        edges = positive_triples[edges]
+        head, rela, tail = np.array(edges).transpose()
+        entity, index = np.unique((head, tail), return_inverse=True) 
+        head, tail = np.reshape(index, (2, -1))
+
+        return np.stack((head,rela,tail)).transpose(), \
+                torch.from_numpy(entity).view(-1,1).long()
+
+    def sampling_negative(
+        self,
+        mode: int,
+        pos_triples: list[tuple[int, int, int]]) -> np.ndarray:
+
+        """采样负三元组。
+
+        :param mode: 'head' 或 'tail'
+        :type mode: str
+        :param pos_triples: 知识图谱中的正确三元组
+        :type pos_triples: list[tuple[int, int, int]]
+        :returns: 负三元组
+        :rtype: numpy.ndarray
+        """
+
+        neg_random = np.random.choice(
+            len(self.entity), 
+            size = self.neg_ent * len(pos_triples)
+        )
+        neg_samples = np.tile(pos_triples, (self.neg_ent, 1))
+        if mode == 'head':
+            neg_samples[:,0] = neg_random
+        elif mode == 'tail':
+            neg_samples[:,2] = neg_random
+        return neg_samples
+
+    def build_graph(
+        self,
+        num_ent: int,
+        triples: tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+        power: int = -1) -> tuple[dgl.DGLGraph, torch.Tensor, torch.Tensor]:
+
+        """建立子图。
+
+        :param num_ent: 子图的节点数
+        :type num_ent: int
+        :param triples: 知识图谱中的正确三元组子集
+        :type triples: tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        :param power: 幂
+        :type power: int
+        :returns: 子图、关系、边的归一化系数
+        :rtype: tuple[dgl.DGLGraph, torch.Tensor, torch.Tensor]
+        """
+
+        head, rela, tail = triples[0], triples[1], triples[2]
+        graph = dgl.graph(([], []))
+        graph.add_nodes(num_ent)
+        graph.add_edges(head, tail)
+        node_norm = self.comp_deg_norm(graph, power)
+        edge_norm = self.node_norm_to_edge_norm(graph,node_norm)
+        rela = torch.tensor(rela)
+        return graph, rela, edge_norm
+
+    def comp_deg_norm(
+        self,
+        graph: dgl.DGLGraph,
+        power: int = -1) -> torch.Tensor:
+
+        """根据目标节点度计算目标节点的归一化系数。
+
+        :param graph: 子图
+        :type graph: dgl.DGLGraph
+        :param power: 幂
+        :type power: int
+        :returns: 节点的归一化系数
+        :rtype: torch.Tensor
+        """
+
+        graph = graph.local_var()
+        in_deg = graph.in_degrees(range(graph.number_of_nodes())).float().numpy()
+        norm = in_deg.__pow__(power)
+        norm[np.isinf(norm)] = 0
+        return torch.from_numpy(norm)
+
+    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 = graph.local_var()
+        # convert to edge norm
+        graph.ndata['norm'] = node_norm.view(-1,1)
+        graph.apply_edges(lambda edges : {'norm' : edges.dst['norm']})
+        return graph.edata['norm']

unike/data/RGCNTestSampler.py

@@ -0,0 +1,208 @@
+# coding:utf-8
+#
+# unike/data/RGCNTestSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 16, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 29, 2024
+#
+# R-GCN 的测试数据采样器.
+
+"""
+RGCNTestSampler - R-GCN 的测试数据采样器。
+"""
+
+import os
+import dgl
+import torch
+import typing
+import numpy as np
+from .TestSampler import TestSampler
+from .RGCNSampler import RGCNSampler
+from .CompGCNSampler import CompGCNSampler
+from typing_extensions import override
+
+class RGCNTestSampler(TestSampler):
+
+    """``R-GCN`` :cite:`R-GCN` 的测试数据采样器。
+
+    例子::
+
+        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: typing.Union[RGCNSampler, CompGCNSampler],
+        valid_file: str = "valid2id.txt",
+        test_file: str = "test2id.txt",
+        type_constrain: bool = True):
+
+        """创建 RGCNTestSampler 对象。
+
+        :param sampler: 训练数据采样器。
+        :type sampler: typing.Union[RGCNSampler, 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.triples: list[tuple[int, int, int]] = self.sampler.t_triples if isinstance(self.sampler, CompGCNSampler) else self.sampler.train_triples
+        #: 幂
+        self.power: float = -1
+
+        self.add_valid_test_reverse_triples()
+        self.get_hr2t_rt2h_from_all()
+
+    @override
+    def get_valid_test_triples_id(self):
+
+        """读取 :py:attr:`valid_file` 文件和 :py:attr:`test_file` 文件。"""
+                
+        with open(os.path.join(self.sampler.in_path, self.valid_file)) as f:
+            self.valid_tol = (int)(f.readline())
+            for line in f:
+                h, t, r = line.strip().split()
+                self.valid_triples.append((int(h), int(r), int(t)))
+                
+        with open(os.path.join(self.sampler.in_path, self.test_file)) as f:
+            self.test_tol = (int)(f.readline())
+            for line in f:
+                h, t, r = line.strip().split()
+                self.test_triples.append((int(h), int(r), int(t)))
+
+    def add_valid_test_reverse_triples(self):
+
+        """对于每一个三元组 (h, r, t)，生成相反关系三元组 (t, r`, h): r` = r + rel_tol。"""
+
+        tol = int(self.sampler.rel_tol / 2)
+                
+        with open(os.path.join(self.sampler.in_path, self.valid_file)) as f:
+            f.readline()
+            for line in f:
+                h, t, r = line.strip().split()
+                self.valid_triples.append(
+                    (int(t), int(r) + tol, int(h))
+                )
+                
+        with open(os.path.join(self.sampler.in_path, self.test_file)) as f:
+            f.readline()
+            for line in f:
+                h, t, r = line.strip().split()
+                self.test_triples.append(
+                    (int(t), int(r) + tol, int(h))
+                )
+                
+        self.all_true_triples = set(
+            self.triples + self.valid_triples + self.test_triples
+        )
+
+    @override
+    def get_type_constrain_id(self):
+
+        """读取 type_constrain.txt 文件。"""
+
+        tol = int(self.sampler.rel_tol / 2)
+                
+        with open(os.path.join(self.sampler.in_path, "type_constrain.txt")) as f:
+            rel_tol = (int)(f.readline())
+            first_line = True
+            for line in f:
+                rel_types = line.strip().split("\t")
+                for entity in rel_types[2:]:
+                    if first_line:
+                        self.rel_heads[int(rel_types[0])].add(int(entity))
+                        self.rel_tails[int(rel_types[0]) + tol].add(int(entity))
+                    else:
+                        self.rel_tails[int(rel_types[0])].add(int(entity))
+                        self.rel_heads[int(rel_types[0]) + tol].add(int(entity))
+                first_line = not first_line
+
+        for rel in self.rel_heads:
+            self.rel_heads[rel] = torch.tensor(list(self.rel_heads[rel]))
+        for rel in self.rel_tails:
+            self.rel_tails[rel] = torch.tensor(list(self.rel_tails[rel]))
+
+    @override
+    def sampling(
+        self,
+        data: list[tuple[int, int, int]]) -> dict[str, typing.Union[dgl.DGLGraph , torch.Tensor]]:
+
+        """``R-GCN`` :cite:`R-GCN` 的测试数据采样函数。
+        
+        :param data: 测试的正确三元组
+        :type data: list[tuple[int, int, int]]
+        :returns: ``R-GCN`` :cite:`R-GCN` 的测试数据
+        :rtype: dict[str, typing.Union[dgl.DGLGraph , torch.Tensor]]
+        """
+        
+        batch_data = {}
+        head_label = torch.zeros(len(data), self.ent_tol)
+        tail_label = torch.zeros(len(data), self.ent_tol)
+        for idx, triple in enumerate(data):
+            head, rel, tail = triple
+            head_label[idx][self.rt2h_all[(rel, tail)]] = 1.0
+            tail_label[idx][self.hr2t_all[(head, rel)]] = 1.0
+        if self.type_constrain:
+            head_label_type = torch.ones(len(data), self.ent_tol)
+            tail_laebl_type = torch.ones(len(data), self.ent_tol)
+            for idx, triple in enumerate(data):
+                head, rel, tail = triple
+                head_label_type[idx][self.rel_heads[rel]] = 0.0
+                tail_laebl_type[idx][self.rel_tails[rel]] = 0.0
+                head_label_type[idx][self.rt2h_all[(rel, tail)]] = 1.0
+                tail_laebl_type[idx][self.hr2t_all[(head, rel)]] = 1.0
+            batch_data["head_label_type"] = head_label_type
+            batch_data["tail_label_type"] = tail_laebl_type
+        batch_data["positive_sample"] = torch.tensor(data)
+        batch_data["head_label"] = head_label
+        batch_data["tail_label"] = tail_label
+        
+        graph, rela, norm = self.sampler.build_graph(self.ent_tol, np.array(self.triples).transpose(), self.power)
+        batch_data["graph"]  = graph
+        batch_data["rela"]   = rela
+        batch_data["norm"]   = norm
+        batch_data["entity"] = torch.arange(0, self.ent_tol, dtype=torch.long).view(-1,1)
+        
+        return batch_data

unike/data/RevSampler.py

@@ -0,0 +1,78 @@
+# coding:utf-8
+#
+# unike/data/RevSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 15, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 16, 2024
+#
+# 为 KGReader 增加相反关系，用于图神经网络模型.
+
+"""
+RevSampler - 为 KGReader 增加相反关系，用于图神经网络模型。
+"""
+
+import os
+from .KGReader import KGReader
+
+class RevSampler(KGReader):
+    
+    """在训练集中增加相反关系.
+
+    对于每一个三元组 (h, r, t)，生成相反关系三元组 (t, r`, h): r` = r + rel_tol。
+    """
+    
+    def __init__(
+        self,
+        in_path: str = "./",
+        ent_file: str = "entity2id.txt",
+        rel_file: str = "relation2id.txt",
+        train_file: str = "train2id.txt"):
+
+        """创建 RevSampler 对象。
+
+        :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
+        """
+        
+        super().__init__(
+            in_path=in_path,
+            ent_file=ent_file,
+            rel_file=rel_file,
+            train_file=train_file
+        )
+
+        self.add_reverse_relation()
+        self.add_train_reverse_triples()
+        self.get_hr2t_rt2h_from_train()
+        
+    def add_reverse_relation(self):
+
+        """增加相反关系：r` = r + rel_tol"""
+        
+        with open(os.path.join(self.in_path, self.rel_file)) as f:
+            f.readline()
+            for line in f:
+                relation, rid = line.strip().split("\t")
+                self.rel2id[relation + "_reverse"] = int(rid) + self.rel_tol
+                self.id2rel[int(rid) + self.rel_tol] = relation + "_reverse"
+        self.rel_tol = len(self.rel2id)
+        
+    def add_train_reverse_triples(self):
+
+        """对于每一个三元组 (h, r, t)，生成相反关系三元组 (t, r`, h): r` = r + rel_tol。"""
+
+        tol = int(self.rel_tol / 2)
+        
+        with open(os.path.join(self.in_path, self.train_file)) as f:
+            f.readline()
+            for line in f:
+                h, t, r = line.strip().split()
+                self.train_triples.append(
+                    (int(t), int(r) + tol, int(h))
+                )