unike 3.0.1__py3-none-any.whl

unike/data/TestSampler.py

@@ -0,0 +1,189 @@
+# coding:utf-8
+#
+# unike/data/TestSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 29, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Feb 25, 2024
+#
+# 测试数据采样器基类.
+
+"""
+TestSampler - 测试数据采样器基类。
+"""
+
+import os
+import torch
+import typing
+from .TradSampler import TradSampler
+from .RGCNSampler import RGCNSampler
+from .CompGCNSampler import CompGCNSampler
+from collections import defaultdict as ddict
+from ..utils import construct_type_constrain
+
+class TestSampler(object):
+
+    """测试数据采样器基类。
+    """
+
+    def __init__(
+        self,
+        sampler: typing.Union[TradSampler, RGCNSampler, CompGCNSampler],
+        valid_file: str = "valid2id.txt",
+        test_file: str = "test2id.txt",
+        type_constrain: bool = True):
+
+        """创建 TestSampler 对象。
+
+        :param sampler: 训练数据采样器。
+        :type sampler: typing.Union[TradSampler, 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
+        """
+
+        #: 训练数据采样器
+        self.sampler: typing.Union[TradSampler, RGCNSampler, CompGCNSampler] = sampler
+        #: 实体的个数
+        self.ent_tol: int = sampler.ent_tol
+        #: valid2id.txt
+        self.valid_file: str = valid_file
+        #: test2id.txt
+        self.test_file: str = test_file
+
+        #: 验证集三元组的个数
+        self.valid_tol: int = 0
+        #: 测试集三元组的个数
+        self.test_tol: int = 0
+
+        #: 验证集三元组
+        self.valid_triples: list[tuple[int, int, int]] = []
+        #: 测试集三元组
+        self.test_triples: list[tuple[int, int, int]] = []
+        #: 知识图谱所有三元组
+        self.all_true_triples: set[tuple[int, int, int]] = set()
+
+        self.get_valid_test_triples_id()
+
+        #: 知识图谱中所有 h-r 对对应的 t 集合
+        self.hr2t_all: ddict[set] = ddict(set)
+        #: 知识图谱中所有 r-t 对对应的 h 集合
+        self.rt2h_all: ddict[set] = ddict(set)
+
+        #: 是否报告 type_constrain.txt 限制的测试结果
+        self.type_constrain: bool = type_constrain
+
+        if self.type_constrain:
+            construct_type_constrain(
+                in_path = self.sampler.in_path,
+                train_file = self.sampler.train_file,
+                valid_file = self.valid_file,
+                test_file = self.test_file
+            )
+            #: 知识图谱中所有 r 存在头实体种类
+            self.rel_heads: ddict[set] = ddict(set)
+            #: 知识图谱中所有 r 存在尾实体种类
+            self.rel_tails: ddict[set] = ddict(set)
+            self.get_type_constrain_id()
+
+    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)))
+
+        self.all_true_triples = set(
+            self.sampler.train_triples + self.valid_triples + self.test_triples
+        )
+        
+    def get_type_constrain_id(self):
+
+        """读取 type_constrain.txt 文件。"""
+                
+        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))
+                    else:
+                        self.rel_tails[int(rel_types[0])].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]))
+
+    def get_hr2t_rt2h_from_all(self):
+
+        """获得 :py:attr:`hr2t_all` 和 :py:attr:`rt2h_all` 。"""
+
+        for h, r, t in self.all_true_triples:
+            self.hr2t_all[(h, r)].add(t)
+            self.rt2h_all[(r, t)].add(h)
+        for h, r in self.hr2t_all:
+            self.hr2t_all[(h, r)] = torch.tensor(list(self.hr2t_all[(h, r)]))
+        for r, t in self.rt2h_all:
+            self.rt2h_all[(r, t)] = torch.tensor(list(self.rt2h_all[(r, t)]))
+
+    def sampling(
+        self,
+        data: list[tuple[int, int, int]]) -> dict[str, torch.Tensor]:
+
+        """采样函数。该方法未实现，子类必须重写该方法，否则抛出 :py:class:`NotImplementedError` 错误。
+        
+        :param data: 测试的正确三元组
+        :type data: list[tuple[int, int, int]]
+        :returns: 测试数据
+        :rtype: dict[str, torch.Tensor]
+        """
+        
+        raise NotImplementedError
+
+    def get_valid(self) -> list[tuple[int, int, int]]:
+
+        """
+        返回验证集三元组。
+
+        :returns: :py:attr:`valid_triples`
+        :rtype: list[tuple[int, int, int]]
+        """
+
+        return self.valid_triples
+
+    def get_test(self) -> list[tuple[int, int, int]]:
+
+        """
+        返回测试集三元组。
+
+        :returns: :py:attr:`test_triples`
+        :rtype: list[tuple[int, int, int]]
+        """
+
+        return self.test_triples
+
+    def get_all_true_triples(self) -> set[tuple[int, int, int]]:
+
+        """
+        返回知识图谱所有三元组。
+
+        :returns: :py:attr:`all_true_triples`
+        :rtype: set[tuple[int, int, int]]
+        """
+
+        return self.all_true_triples 

unike/data/TradSampler.py

@@ -0,0 +1,122 @@
+# coding:utf-8
+#
+# unike/data/TradSampler.py
+#
+# created by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 28, 2024
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Feb 25, 2024
+#
+# 为 KGReader 增加构建负三元组的函数，用于平移模型和语义匹配模型.
+
+"""
+TradSampler - 为 KGReader 增加构建负三元组的函数，用于平移模型和语义匹配模型。
+"""
+
+import torch
+import typing
+import numpy as np
+from .KGReader import KGReader
+
+class TradSampler(KGReader):
+    
+    """
+    平移模型和语义匹配模型的采样器的基类。
+    """
+    
+    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):
+
+        """创建 TradSampler 对象。
+
+        :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
+        self.batch_size: int = batch_size
+        #: 对于每一个正三元组, 构建的负三元组的个数, 替换 entity (head + tail)
+        self.neg_ent: int = neg_ent
+
+        self.get_hr2t_rt2h_from_train()
+
+    def sampling(
+        self,
+        pos_triples: list[tuple[int, int, int]]) -> dict[str, typing.Union[str, torch.Tensor]]:
+        
+        """平移模型和语义匹配模型的训练集数据采样函数。该方法未实现，子类必须重写该方法，否则抛出 :py:class:`NotImplementedError` 错误。
+        
+        :param pos_triples: 知识图谱中的正确三元组
+        :type pos_triples: list[tuple[int, int, int]]
+        :returns: 平移模型和语义匹配模型的训练数据
+        :rtype: dict[str, typing.Union[str, torch.Tensor]]
+        """
+        
+        raise NotImplementedError
+
+    def corrupt_head(
+        self,
+        t: int,
+        r: int,
+        num_max: int = 1) -> np.ndarray:
+
+        """替换头实体构建负三元组。
+
+        :param t: 尾实体
+        :type t: int
+        :param r: 关系
+        :type r: int
+        :param num_max: 一次负采样的个数
+        :type num_max: int
+        :returns: 负三元组的头实体列表
+        :rtype: numpy.ndarray
+        """
+        
+        tmp = torch.randint(low=0, high=self.ent_tol, size=(num_max,)).numpy()
+        mask = np.in1d(tmp, self.rt2h_train[(r, t)], assume_unique=True, invert=True)
+        neg = tmp[mask]
+        return neg
+        
+    def corrupt_tail(
+        self,
+        h: int,
+        r: int,
+        num_max: int = 1) -> np.ndarray:
+        
+        """替换尾实体构建负三元组。
+
+        :param h: 头实体
+        :type h: int
+        :param r: 关系
+        :type r: int
+        :param num_max: 一次负采样的个数
+        :type num_max: int
+        :returns: 负三元组的尾实体列表
+        :rtype: numpy.ndarray
+        """
+        
+        tmp = torch.randint(low=0, high=self.ent_tol, size=(num_max,)).numpy()
+        mask = np.in1d(tmp, self.hr2t_train[(h, r)], assume_unique=True, invert=True)
+        neg = tmp[mask]
+        return neg

unike/data/TradTestSampler.py

@@ -0,0 +1,87 @@
+# coding:utf-8
+#
+# unike/data/TradTestSampler.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
+#
+# 平移模型和语义匹配模型的测试数据采样器.
+
+"""
+TradTestSampler - 平移模型和语义匹配模型的测试数据采样器。
+"""
+
+import torch
+from .TradSampler import TradSampler
+from .TestSampler import TestSampler
+from typing_extensions import override
+
+class TradTestSampler(TestSampler):
+
+    """平移模型和语义匹配模型的测试数据采样器。
+    """
+
+    def __init__(
+        self,
+        sampler: TradSampler,
+        valid_file: str = "valid2id.txt",
+        test_file: str = "test2id.txt",
+        type_constrain: bool = True):
+
+        """创建 TradTestSampler 对象。
+
+        :param sampler: 训练数据采样器。
+        :type sampler: TradSampler
+        :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.get_hr2t_rt2h_from_all()
+
+    @override
+    def sampling(
+        self,
+        data: list[tuple[int, int, int]]) -> dict[str, torch.Tensor]:
+
+        """采样函数。
+        
+        :param data: 测试的正确三元组
+        :type data: list[tuple[int, int, int]]
+        :returns: 测试数据
+        :rtype: dict[str, 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
+        return batch_data

unike/data/UniSampler.py

@@ -0,0 +1,145 @@
+# coding:utf-8
+#
+# unike/data/UniSampler.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
+#
+# 平移模型和语义匹配模型的训练集数据采样器.
+
+"""
+UniSampler - 平移模型和语义匹配模型的训练集数据采样器。
+"""
+
+import torch
+import typing
+import warnings
+import numpy as np
+from .TradSampler import TradSampler
+from typing_extensions import override
+
+warnings.filterwarnings("ignore")
+
+class UniSampler(TradSampler):
+
+    """平移模型和语义匹配模型的训练集普通的数据采样器（均值分布）。
+    """
+    
+    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):
+
+        """创建 UniSampler 对象。
+
+        :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.cross_sampling_flag = 0
+
+    @override
+    def sampling(
+        self,
+        pos_triples: list[tuple[int, int, int]]) -> dict[str, typing.Union[str, torch.Tensor]]:
+        
+        """平移模型和语义匹配模型的训练集普通的数据采样函数（均匀分布）。
+        
+        :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 = []
+        self.cross_sampling_flag = 1 - self.cross_sampling_flag
+        if self.cross_sampling_flag == 0:
+            batch_data['mode'] = "head-batch"
+            for h, r, t in pos_triples:
+                neg_head = self.head_batch(t, r, self.neg_ent)
+                neg_ent_sample.append(neg_head)
+        else:
+            batch_data['mode'] = "tail-batch"
+            for h, r, t in pos_triples:
+                neg_tail = self.tail_batch(h, r, self.neg_ent)
+                neg_ent_sample.append(neg_tail)
+
+        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 head_batch(
+        self,
+        t: int,
+        r: int,
+        neg_size: int= None) -> np.ndarray:
+
+        """替换头实体构建负三元组。
+
+        :param t: 尾实体
+        :type t: int
+        :param r: 关系
+        :type r: int
+        :param neg_size: 负三元组个数
+        :type neg_size: int
+        :returns: 负三元组中的头实体列表
+        :rtype: numpy.ndarray
+        """
+        
+        neg_list = []
+        neg_cur_size = 0
+        while neg_cur_size < neg_size:
+            neg_tmp = self.corrupt_head(t, r, num_max=(neg_size - neg_cur_size) * 2)
+            neg_list.append(neg_tmp)
+            neg_cur_size += len(neg_tmp)
+        return np.concatenate(neg_list)[:neg_size]
+
+    def tail_batch(
+        self,
+        h: int,
+        r: int,
+        neg_size: int = None) -> np.ndarray:
+        
+        """替换尾实体构建负三元组。
+
+        :param h: 头实体
+        :type h: int
+        :param r: 关系
+        :type r: int
+        :param neg_size: 负三元组个数
+        :type neg_size: int
+        :returns: 负三元组中的尾实体列表
+        :rtype: numpy.ndarray
+        """
+        
+        neg_list = []
+        neg_cur_size = 0
+        while neg_cur_size < neg_size:
+            neg_tmp = self.corrupt_tail(h, r, num_max=(neg_size - neg_cur_size) * 2)
+            neg_list.append(neg_tmp)
+            neg_cur_size += len(neg_tmp)
+        return np.concatenate(neg_list)[:neg_size]

unike/data/__init__.py

@@ -0,0 +1,47 @@
+# coding:utf-8
+#
+# unike/data/__init__.py
+# 
+# git pull from OpenKE-PyTorch by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on May 7, 2023
+# updated by LuYF-Lemon-love <luyanfeng_nlp@qq.com> on Jan 29, 2023
+# 
+# 该头文件定义了 data 接口.
+
+"""数据采样部分，包含为训练和验证模型定义的数据采样器。"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .KGReader import KGReader
+
+from .TradSampler import TradSampler
+from .UniSampler import UniSampler
+from .BernSampler import BernSampler
+
+from .RevSampler import RevSampler
+from .RGCNSampler import RGCNSampler
+from .CompGCNSampler import CompGCNSampler
+
+from .TestSampler import TestSampler
+from .TradTestSampler import TradTestSampler
+from .RGCNTestSampler import RGCNTestSampler
+from .CompGCNTestSampler import CompGCNTestSampler
+
+from .KGEDataLoader import KGEDataLoader, get_kge_data_loader_hpo_config
+
+__all__ = [
+	'KGReader',
+	'TradSampler',
+	'UniSampler',
+	'BernSampler',
+	'RevSampler',
+	'RGCNSampler',
+	'CompGCNSampler',
+	'TestSampler',
+	'TradTestSampler',
+	'RGCNTestSampler',
+	'CompGCNTestSampler',
+	'KGEDataLoader',
+	'get_kge_data_loader_hpo_config'
+]