├── LICENSE.txt ├── README.md └── TransE_type ├── TransE_type.py ├── dataset.py └── train_transe_type.py /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Attribution-NonCommercial-ShareAlike 4.0 International 2 | 3 | ======================================================================= 4 | 5 | Creative Commons Corporation ("Creative Commons") is not a law firm and 6 | does not provide legal services or legal advice. Distribution of 7 | Creative Commons public licenses does not create a lawyer-client or 8 | other relationship. Creative Commons makes its licenses and related 9 | information available on an "as-is" basis. Creative Commons gives no 10 | warranties regarding its licenses, any material licensed under their 11 | terms and conditions, or any related information. Creative Commons 12 | disclaims all liability for damages resulting from their use to the 13 | fullest extent possible. 14 | 15 | Using Creative Commons Public Licenses 16 | 17 | Creative Commons public licenses provide a standard set of terms and 18 | conditions that creators and other rights holders may use to share 19 | original works of authorship and other material subject to copyright 20 | and certain other rights specified in the public license below. The 21 | following considerations are for informational purposes only, are not 22 | exhaustive, and do not form part of our licenses. 23 | 24 | Considerations for licensors: Our public licenses are 25 | intended for use by those authorized to give the public 26 | permission to use material in ways otherwise restricted by 27 | copyright and certain other rights. Our licenses are 28 | irrevocable. Licensors should read and understand the terms 29 | and conditions of the license they choose before applying it. 30 | Licensors should also secure all rights necessary before 31 | applying our licenses so that the public can reuse the 32 | material as expected. Licensors should clearly mark any 33 | material not subject to the license. This includes other CC- 34 | licensed material, or material used under an exception or 35 | limitation to copyright. More considerations for licensors: 36 | wiki.creativecommons.org/Considerations_for_licensors 37 | 38 | Considerations for the public: By using one of our public 39 | licenses, a licensor grants the public permission to use the 40 | licensed material under specified terms and conditions. If 41 | the licensor's permission is not necessary for any reason--for 42 | example, because of any applicable exception or limitation to 43 | copyright--then that use is not regulated by the license. Our 44 | licenses grant only permissions under copyright and certain 45 | other rights that a licensor has authority to grant. Use of 46 | the licensed material may still be restricted for other 47 | reasons, including because others have copyright or other 48 | rights in the material. A licensor may make special requests, 49 | such as asking that all changes be marked or described. 50 | Although not required by our licenses, you are encouraged to 51 | respect those requests where reasonable. More considerations 52 | for the public: 53 | wiki.creativecommons.org/Considerations_for_licensees 54 | 55 | ======================================================================= 56 | 57 | Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International 58 | Public License 59 | 60 | By exercising the Licensed Rights (defined below), You accept and agree 61 | to be bound by the terms and conditions of this Creative Commons 62 | Attribution-NonCommercial-ShareAlike 4.0 International Public License 63 | ("Public License"). To the extent this Public License may be 64 | interpreted as a contract, You are granted the Licensed Rights in 65 | consideration of Your acceptance of these terms and conditions, and the 66 | Licensor grants You such rights in consideration of benefits the 67 | Licensor receives from making the Licensed Material available under 68 | these terms and conditions. 69 | 70 | 71 | Section 1 -- Definitions. 72 | 73 | a. Adapted Material means material subject to Copyright and Similar 74 | Rights that is derived from or based upon the Licensed Material 75 | and in which the Licensed Material is translated, altered, 76 | arranged, transformed, or otherwise modified in a manner requiring 77 | permission under the Copyright and Similar Rights held by the 78 | Licensor. For purposes of this Public License, where the Licensed 79 | Material is a musical work, performance, or sound recording, 80 | Adapted Material is always produced where the Licensed Material is 81 | synched in timed relation with a moving image. 82 | 83 | b. Adapter's License means the license You apply to Your Copyright 84 | and Similar Rights in Your contributions to Adapted Material in 85 | accordance with the terms and conditions of this Public License. 86 | 87 | c. BY-NC-SA Compatible License means a license listed at 88 | creativecommons.org/compatiblelicenses, approved by Creative 89 | Commons as essentially the equivalent of this Public License. 90 | 91 | d. Copyright and Similar Rights means copyright and/or similar rights 92 | closely related to copyright including, without limitation, 93 | performance, broadcast, sound recording, and Sui Generis Database 94 | Rights, without regard to how the rights are labeled or 95 | categorized. For purposes of this Public License, the rights 96 | specified in Section 2(b)(1)-(2) are not Copyright and Similar 97 | Rights. 98 | 99 | e. Effective Technological Measures means those measures that, in the 100 | absence of proper authority, may not be circumvented under laws 101 | fulfilling obligations under Article 11 of the WIPO Copyright 102 | Treaty adopted on December 20, 1996, and/or similar international 103 | agreements. 104 | 105 | f. Exceptions and Limitations means fair use, fair dealing, and/or 106 | any other exception or limitation to Copyright and Similar Rights 107 | that applies to Your use of the Licensed Material. 108 | 109 | g. License Elements means the license attributes listed in the name 110 | of a Creative Commons Public License. The License Elements of this 111 | Public License are Attribution, NonCommercial, and ShareAlike. 112 | 113 | h. Licensed Material means the artistic or literary work, database, 114 | or other material to which the Licensor applied this Public 115 | License. 116 | 117 | i. Licensed Rights means the rights granted to You subject to the 118 | terms and conditions of this Public License, which are limited to 119 | all Copyright and Similar Rights that apply to Your use of the 120 | Licensed Material and that the Licensor has authority to license. 121 | 122 | j. Licensor means the individual(s) or entity(ies) granting rights 123 | under this Public License. 124 | 125 | k. NonCommercial means not primarily intended for or directed towards 126 | commercial advantage or monetary compensation. For purposes of 127 | this Public License, the exchange of the Licensed Material for 128 | other material subject to Copyright and Similar Rights by digital 129 | file-sharing or similar means is NonCommercial provided there is 130 | no payment of monetary compensation in connection with the 131 | exchange. 132 | 133 | l. Share means to provide material to the public by any means or 134 | process that requires permission under the Licensed Rights, such 135 | as reproduction, public display, public performance, distribution, 136 | dissemination, communication, or importation, and to make material 137 | available to the public including in ways that members of the 138 | public may access the material from a place and at a time 139 | individually chosen by them. 140 | 141 | m. Sui Generis Database Rights means rights other than copyright 142 | resulting from Directive 96/9/EC of the European Parliament and of 143 | the Council of 11 March 1996 on the legal protection of databases, 144 | as amended and/or succeeded, as well as other essentially 145 | equivalent rights anywhere in the world. 146 | 147 | n. You means the individual or entity exercising the Licensed Rights 148 | under this Public License. Your has a corresponding meaning. 149 | 150 | 151 | Section 2 -- Scope. 152 | 153 | a. License grant. 154 | 155 | 1. Subject to the terms and conditions of this Public License, 156 | the Licensor hereby grants You a worldwide, royalty-free, 157 | non-sublicensable, non-exclusive, irrevocable license to 158 | exercise the Licensed Rights in the Licensed Material to: 159 | 160 | a. reproduce and Share the Licensed Material, in whole or 161 | in part, for NonCommercial purposes only; and 162 | 163 | b. produce, reproduce, and Share Adapted Material for 164 | NonCommercial purposes only. 165 | 166 | 2. Exceptions and Limitations. For the avoidance of doubt, where 167 | Exceptions and Limitations apply to Your use, this Public 168 | License does not apply, and You do not need to comply with 169 | its terms and conditions. 170 | 171 | 3. Term. The term of this Public License is specified in Section 172 | 6(a). 173 | 174 | 4. Media and formats; technical modifications allowed. The 175 | Licensor authorizes You to exercise the Licensed Rights in 176 | all media and formats whether now known or hereafter created, 177 | and to make technical modifications necessary to do so. The 178 | Licensor waives and/or agrees not to assert any right or 179 | authority to forbid You from making technical modifications 180 | necessary to exercise the Licensed Rights, including 181 | technical modifications necessary to circumvent Effective 182 | Technological Measures. For purposes of this Public License, 183 | simply making modifications authorized by this Section 2(a) 184 | (4) never produces Adapted Material. 185 | 186 | 5. Downstream recipients. 187 | 188 | a. Offer from the Licensor -- Licensed Material. Every 189 | recipient of the Licensed Material automatically 190 | receives an offer from the Licensor to exercise the 191 | Licensed Rights under the terms and conditions of this 192 | Public License. 193 | 194 | b. Additional offer from the Licensor -- Adapted Material. 195 | Every recipient of Adapted Material from You 196 | automatically receives an offer from the Licensor to 197 | exercise the Licensed Rights in the Adapted Material 198 | under the conditions of the Adapter's License You apply. 199 | 200 | c. No downstream restrictions. You may not offer or impose 201 | any additional or different terms or conditions on, or 202 | apply any Effective Technological Measures to, the 203 | Licensed Material if doing so restricts exercise of the 204 | Licensed Rights by any recipient of the Licensed 205 | Material. 206 | 207 | 6. No endorsement. Nothing in this Public License constitutes or 208 | may be construed as permission to assert or imply that You 209 | are, or that Your use of the Licensed Material is, connected 210 | with, or sponsored, endorsed, or granted official status by, 211 | the Licensor or others designated to receive attribution as 212 | provided in Section 3(a)(1)(A)(i). 213 | 214 | b. Other rights. 215 | 216 | 1. Moral rights, such as the right of integrity, are not 217 | licensed under this Public License, nor are publicity, 218 | privacy, and/or other similar personality rights; however, to 219 | the extent possible, the Licensor waives and/or agrees not to 220 | assert any such rights held by the Licensor to the limited 221 | extent necessary to allow You to exercise the Licensed 222 | Rights, but not otherwise. 223 | 224 | 2. Patent and trademark rights are not licensed under this 225 | Public License. 226 | 227 | 3. To the extent possible, the Licensor waives any right to 228 | collect royalties from You for the exercise of the Licensed 229 | Rights, whether directly or through a collecting society 230 | under any voluntary or waivable statutory or compulsory 231 | licensing scheme. In all other cases the Licensor expressly 232 | reserves any right to collect such royalties, including when 233 | the Licensed Material is used other than for NonCommercial 234 | purposes. 235 | 236 | 237 | Section 3 -- License Conditions. 238 | 239 | Your exercise of the Licensed Rights is expressly made subject to the 240 | following conditions. 241 | 242 | a. Attribution. 243 | 244 | 1. If You Share the Licensed Material (including in modified 245 | form), You must: 246 | 247 | a. retain the following if it is supplied by the Licensor 248 | with the Licensed Material: 249 | 250 | i. identification of the creator(s) of the Licensed 251 | Material and any others designated to receive 252 | attribution, in any reasonable manner requested by 253 | the Licensor (including by pseudonym if 254 | designated); 255 | 256 | ii. a copyright notice; 257 | 258 | iii. a notice that refers to this Public License; 259 | 260 | iv. a notice that refers to the disclaimer of 261 | warranties; 262 | 263 | v. a URI or hyperlink to the Licensed Material to the 264 | extent reasonably practicable; 265 | 266 | b. indicate if You modified the Licensed Material and 267 | retain an indication of any previous modifications; and 268 | 269 | c. indicate the Licensed Material is licensed under this 270 | Public License, and include the text of, or the URI or 271 | hyperlink to, this Public License. 272 | 273 | 2. You may satisfy the conditions in Section 3(a)(1) in any 274 | reasonable manner based on the medium, means, and context in 275 | which You Share the Licensed Material. For example, it may be 276 | reasonable to satisfy the conditions by providing a URI or 277 | hyperlink to a resource that includes the required 278 | information. 279 | 3. If requested by the Licensor, You must remove any of the 280 | information required by Section 3(a)(1)(A) to the extent 281 | reasonably practicable. 282 | 283 | b. ShareAlike. 284 | 285 | In addition to the conditions in Section 3(a), if You Share 286 | Adapted Material You produce, the following conditions also apply. 287 | 288 | 1. The Adapter's License You apply must be a Creative Commons 289 | license with the same License Elements, this version or 290 | later, or a BY-NC-SA Compatible License. 291 | 292 | 2. You must include the text of, or the URI or hyperlink to, the 293 | Adapter's License You apply. You may satisfy this condition 294 | in any reasonable manner based on the medium, means, and 295 | context in which You Share Adapted Material. 296 | 297 | 3. You may not offer or impose any additional or different terms 298 | or conditions on, or apply any Effective Technological 299 | Measures to, Adapted Material that restrict exercise of the 300 | rights granted under the Adapter's License You apply. 301 | 302 | 303 | Section 4 -- Sui Generis Database Rights. 304 | 305 | Where the Licensed Rights include Sui Generis Database Rights that 306 | apply to Your use of the Licensed Material: 307 | 308 | a. for the avoidance of doubt, Section 2(a)(1) grants You the right 309 | to extract, reuse, reproduce, and Share all or a substantial 310 | portion of the contents of the database for NonCommercial purposes 311 | only; 312 | 313 | b. if You include all or a substantial portion of the database 314 | contents in a database in which You have Sui Generis Database 315 | Rights, then the database in which You have Sui Generis Database 316 | Rights (but not its individual contents) is Adapted Material, 317 | including for purposes of Section 3(b); and 318 | 319 | c. You must comply with the conditions in Section 3(a) if You Share 320 | all or a substantial portion of the contents of the database. 321 | 322 | For the avoidance of doubt, this Section 4 supplements and does not 323 | replace Your obligations under this Public License where the Licensed 324 | Rights include other Copyright and Similar Rights. 325 | 326 | 327 | Section 5 -- Disclaimer of Warranties and Limitation of Liability. 328 | 329 | a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE 330 | EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS 331 | AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF 332 | ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS, 333 | IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION, 334 | WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR 335 | PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS, 336 | ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT 337 | KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT 338 | ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU. 339 | 340 | b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE 341 | TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION, 342 | NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT, 343 | INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES, 344 | COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR 345 | USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN 346 | ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR 347 | DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR 348 | IN PART, THIS LIMITATION MAY NOT APPLY TO YOU. 349 | 350 | c. The disclaimer of warranties and limitation of liability provided 351 | above shall be interpreted in a manner that, to the extent 352 | possible, most closely approximates an absolute disclaimer and 353 | waiver of all liability. 354 | 355 | 356 | Section 6 -- Term and Termination. 357 | 358 | a. This Public License applies for the term of the Copyright and 359 | Similar Rights licensed here. However, if You fail to comply with 360 | this Public License, then Your rights under this Public License 361 | terminate automatically. 362 | 363 | b. Where Your right to use the Licensed Material has terminated under 364 | Section 6(a), it reinstates: 365 | 366 | 1. automatically as of the date the violation is cured, provided 367 | it is cured within 30 days of Your discovery of the 368 | violation; or 369 | 370 | 2. upon express reinstatement by the Licensor. 371 | 372 | For the avoidance of doubt, this Section 6(b) does not affect any 373 | right the Licensor may have to seek remedies for Your violations 374 | of this Public License. 375 | 376 | c. For the avoidance of doubt, the Licensor may also offer the 377 | Licensed Material under separate terms or conditions or stop 378 | distributing the Licensed Material at any time; however, doing so 379 | will not terminate this Public License. 380 | 381 | d. Sections 1, 5, 6, 7, and 8 survive termination of this Public 382 | License. 383 | 384 | 385 | Section 7 -- Other Terms and Conditions. 386 | 387 | a. The Licensor shall not be bound by any additional or different 388 | terms or conditions communicated by You unless expressly agreed. 389 | 390 | b. Any arrangements, understandings, or agreements regarding the 391 | Licensed Material not stated herein are separate from and 392 | independent of the terms and conditions of this Public License. 393 | 394 | 395 | Section 8 -- Interpretation. 396 | 397 | a. For the avoidance of doubt, this Public License does not, and 398 | shall not be interpreted to, reduce, limit, restrict, or impose 399 | conditions on any use of the Licensed Material that could lawfully 400 | be made without permission under this Public License. 401 | 402 | b. To the extent possible, if any provision of this Public License is 403 | deemed unenforceable, it shall be automatically reformed to the 404 | minimum extent necessary to make it enforceable. If the provision 405 | cannot be reformed, it shall be severed from this Public License 406 | without affecting the enforceability of the remaining terms and 407 | conditions. 408 | 409 | c. No term or condition of this Public License will be waived and no 410 | failure to comply consented to unless expressly agreed to by the 411 | Licensor. 412 | 413 | d. Nothing in this Public License constitutes or may be interpreted 414 | as a limitation upon, or waiver of, any privileges and immunities 415 | that apply to the Licensor or You, including from the legal 416 | processes of any jurisdiction or authority. 417 | 418 | ======================================================================= 419 | 420 | Creative Commons is not a party to its public 421 | licenses. Notwithstanding, Creative Commons may elect to apply one of 422 | its public licenses to material it publishes and in those instances 423 | will be considered the “Licensor.” The text of the Creative Commons 424 | public licenses is dedicated to the public domain under the CC0 Public 425 | Domain Dedication. Except for the limited purpose of indicating that 426 | material is shared under a Creative Commons public license or as 427 | otherwise permitted by the Creative Commons policies published at 428 | creativecommons.org/policies, Creative Commons does not authorize the 429 | use of the trademark "Creative Commons" or any other trademark or logo 430 | of Creative Commons without its prior written consent including, 431 | without limitation, in connection with any unauthorized modifications 432 | to any of its public licenses or any other arrangements, 433 | understandings, or agreements concerning use of licensed material. For 434 | the avoidance of doubt, this paragraph does not form part of the 435 | public licenses. 436 | 437 | Creative Commons may be contacted at creativecommons.org. 438 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Copyright: KAIST Big Data Intelligence Lab (https://bdi-lab.kaist.ac.kr/) 2 | 3 | Knowledge Graph Embedding with Entity Type Constraints, Seunghwan Kong, Chanyoung Chung, Suheon Ju, and Joyce Jiyoung Whang, Journal of KIISE, Sep. 2022 4 | 5 | All codes are written by Seunghwan Kong ([shkong@kaist.ac.kr](mailto:shkong@kaist.ac.kr)) and proofread by Chanyoung Chung ([chanyoung.chung@kaist.ac.kr](mailto:chanyoung.chung@kaist.ac.kr)). 6 | 7 | ### Data 8 | 9 | To run the code, the dataset should be located in "./data/[DATASET_NAME]" directory with the following files: 10 | 11 | - entity2id.txt: Each line is in the form of "entity entity_index". The first line indicates the number of entities. 12 | - entity2typeid.txt: Each line is in the form of "entity_index type_index". 13 | - relation2id.txt: Each line is in the form of "relation relation_index". The first line indicates the number of relations. 14 | - train2id.txt / valid2id.txt / test2id.txt: Each line is in the form of "head_entity_index, tail_entity_index, relation_index". First line indicates the number of triplets. 15 | - triplets.txt: Each line is in the form of "head_entity relation tail_entity". The first line indicates the number of triplets. 16 | - type2id.txt: Each line is in the form of "type type_index". The first line indicates the number of types. 17 | 18 | ### train_transe_type.py 19 | 20 | Code for generating knowledge graph embeddings with entity type constraints. 21 | 22 | **Arguments** 23 | 24 | - neg: Number of negative samples to generate for each triplet (25 as default) 25 | - dim: Dimension of the embeddings (128 as default) 26 | - epochs: Total epochs for experiment (1000 as default) 27 | - valid_epochs: Number of epochs between validations (50 as default) 28 | - lamb: Weight of type loss (0.01 as default) 29 | - data: Name of the dataset (None as default) 30 | - lr: Learning rate (2.0 as default) 31 | - margin: Margin of the margin ranking loss (1.0 as default) 32 | - test: 0 for the validation 1 for the test (0 as default) 33 | 34 | **Usage:** python3 train_transe_type.py [arguments] 35 | 36 | - Example: python3 train_transe_type.py -lamb 0.05 -data [DATASET_NAME] -lr 0.1 37 | 38 | **Output** 39 | 40 | - Mean rank, mean reciprocal rank, hit@1,3,10 of the trained model in every validation epoch 41 | 42 | ### License 43 | 44 | Our codes are released under the CC BY-NC-SA 4.0 license. 45 | -------------------------------------------------------------------------------- /TransE_type/TransE_type.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from torch.autograd import Variable 5 | import numpy as np 6 | 7 | class TransE_type(nn.Module): 8 | def __init__(self, num_ent, num_rel, dim, p_norm, cuda=True): 9 | super(TransE_type, self).__init__() 10 | 11 | self.num_ent = num_ent 12 | self.num_rel = num_rel 13 | self.dim = dim 14 | self.p_norm = p_norm 15 | 16 | self.ent_embeddings = nn.Embedding(self.num_ent, self.dim) 17 | self.rel_embeddings = nn.Embedding(self.num_rel, self.dim) 18 | 19 | init_range = 6.0 / np.sqrt(self.dim) 20 | nn.init.uniform_(self.ent_embeddings.weight.data, -init_range, init_range) 21 | nn.init.uniform_(self.rel_embeddings.weight.data, -init_range, init_range) 22 | 23 | def forward(self, batch): 24 | h = self.ent_embeddings(batch[:, 0]) 25 | t = self.ent_embeddings(batch[:, 1]) 26 | r = self.rel_embeddings(batch[:, 2]) 27 | 28 | h = F.normalize(h, self.p_norm, -1) 29 | t = F.normalize(t, self.p_norm, -1) 30 | r = F.normalize(r, self.p_norm, -1) 31 | 32 | res = h + r - t 33 | return -torch.norm(res, self.p_norm, -1) -------------------------------------------------------------------------------- /TransE_type/dataset.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from torch.utils.data import Dataset 3 | 4 | class TrainData(Dataset): 5 | def __init__(self, path_data, num_neg): 6 | self.path_data = path_data 7 | self.num_ent = self.num_instance("entity2id.txt") 8 | self.num_rel = self.num_instance("relation2id.txt") 9 | self.num_train, self.train = self.read_triplet("train2id.txt") 10 | self.num_type = self.num_instance("type2id.txt") 11 | self.type2id, self.id2type = self.read_type("entity2typeid.txt") 12 | 13 | self.filter_head, self.filter_tail = self.build_filter_dict() 14 | self.num_neg = num_neg 15 | 16 | def __len__(self): 17 | return self.num_train 18 | 19 | def __getitem__(self, idx): 20 | pos = self.train[idx] 21 | neg = self.negative_sampling(pos) 22 | return pos, neg 23 | 24 | def num_instance(self, target): 25 | with open(self.path_data + target, 'r') as f: 26 | res = int(f.readline().strip()) 27 | return res 28 | 29 | def read_triplet(self, target): 30 | list_tri = [] 31 | with open(self.path_data + target, 'r') as f: 32 | num_tri = int(f.readline().strip()) 33 | for line in f.readlines(): 34 | h, t, r = line.strip().split(' ') 35 | list_tri.append([int(h), int(t), int(r)]) 36 | list_tri = np.array(list_tri) 37 | return num_tri, list_tri 38 | 39 | def read_type(self, target): 40 | type2id = dict() 41 | id2type = [] 42 | with open(self.path_data + target, 'r') as f: 43 | for line in f.readlines(): 44 | idx, typeid = line.strip().split(' ') 45 | idx = int(idx) 46 | typeid = int(typeid) 47 | 48 | if typeid not in type2id: 49 | type2id[typeid] = [] 50 | type2id[typeid].append(idx) 51 | id2type.append(typeid) 52 | return type2id, id2type 53 | 54 | def build_filter_dict(self): 55 | filter_head = dict() 56 | filter_tail = dict() 57 | for h, t, r in self.train: 58 | if (h, r) not in filter_tail: 59 | filter_tail[(h, r)] = [] 60 | filter_tail[(h, r)].append(t) 61 | 62 | if (t, r) not in filter_head: 63 | filter_head[(t, r)] = [] 64 | filter_head[(t, r)].append(h) 65 | 66 | for item in filter_head: 67 | filter_head[item] = np.array(filter_head[item]) 68 | for item in filter_tail: 69 | filter_tail[item] = np.array(filter_tail[item]) 70 | 71 | return filter_head, filter_tail 72 | 73 | def negative_sampling(self, triplet): 74 | h, t, r = triplet 75 | type_h = self.id2type[h] 76 | type_t = self.id2type[t] 77 | count_h = len(self.type2id[type_h]) 78 | count_t = len(self.type2id[type_t]) 79 | 80 | candidate_head = np.setdiff1d(self.type2id[type_h], self.filter_head[(t, r)]) 81 | candidate_tail = np.setdiff1d(self.type2id[type_t], self.filter_tail[(h, r)]) 82 | corrupt = np.random.rand(self.num_neg) 83 | 84 | res = np.tile(triplet, [self.num_neg, 1]) 85 | corrupt_head = np.where(corrupt < count_h / (count_h + count_t))[0] 86 | corrupt_tail = np.where(corrupt >= count_h / (count_h + count_t))[0] 87 | 88 | res[corrupt_head, 0] = np.random.choice(candidate_head, len(corrupt_head)) 89 | res[corrupt_tail, 1] = np.random.choice(candidate_tail, len(corrupt_tail)) 90 | return res 91 | 92 | class ValidData(Dataset): 93 | def __init__(self, path_data, test): 94 | self.path_data = path_data 95 | self.test = test 96 | self.num_ent = self.num_instance("entity2id.txt") 97 | self.num_rel = self.num_instance("relation2id.txt") 98 | self.num_train, self.train = self.read_triplet("train2id.txt") 99 | self.num_valid, self.valid = self.read_triplet("valid2id.txt") 100 | self.num_test, self.test = self.read_triplet("test2id.txt") 101 | self.num_type = self.num_instance("type2id.txt") 102 | self.type2id, self.id2type = self.read_type("entity2typeid.txt") 103 | 104 | self.filter_head, self.filter_tail = self.build_filter_dict() 105 | 106 | def __len__(self): 107 | if self.test: 108 | return self.num_test 109 | else: 110 | return self.num_valid 111 | 112 | def __getitem__(self, idx): 113 | if self.test: 114 | return self.test[idx] 115 | else: 116 | return self.valid[idx] 117 | 118 | def num_instance(self, target): 119 | with open(self.path_data + target, 'r') as f: 120 | res = int(f.readline().strip()) 121 | return res 122 | 123 | def read_triplet(self, target): 124 | list_tri = [] 125 | with open(self.path_data + target, 'r') as f: 126 | num_tri = int(f.readline().strip()) 127 | for line in f.readlines(): 128 | h, t, r = line.strip().split(' ') 129 | list_tri.append([int(h), int(t), int(r)]) 130 | list_tri = np.array(list_tri) 131 | return num_tri, list_tri 132 | 133 | def read_type(self, target): 134 | type2id = dict() 135 | id2type = [] 136 | with open(self.path_data + target, 'r') as f: 137 | for line in f.readlines(): 138 | idx, typeid = line.strip().split(' ') 139 | idx = int(idx) 140 | typeid = int(typeid) 141 | 142 | if typeid not in type2id: 143 | type2id[typeid] = [] 144 | type2id[typeid].append(idx) 145 | id2type.append(typeid) 146 | return type2id, id2type 147 | 148 | def build_filter_dict(self): 149 | filter_head = dict() 150 | filter_tail = dict() 151 | for data in [self.train, self.valid, self.test]: 152 | for h, t, r in data: 153 | if (h, r) not in filter_tail: 154 | filter_tail[(h, r)] = [] 155 | filter_tail[(h, r)].append(t) 156 | 157 | if (t, r) not in filter_head: 158 | filter_head[(t, r)] = [] 159 | filter_head[(t, r)].append(h) 160 | 161 | for item in filter_head: 162 | filter_head[item] = np.array(filter_head[item]) 163 | for item in filter_tail: 164 | filter_tail[item] = np.array(filter_tail[item]) 165 | 166 | return filter_head, filter_tail -------------------------------------------------------------------------------- /TransE_type/train_transe_type.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | from dataset import TrainData, ValidData 5 | from TransE_type import TransE_type 6 | from torch.utils.data import DataLoader 7 | from tqdm import tqdm 8 | import numpy as np 9 | import argparse 10 | import random 11 | 12 | 13 | if __name__ =="__main__": 14 | parser = argparse.ArgumentParser(description='Arguments') 15 | parser.add_argument('--neg', type = int, default = 25) 16 | parser.add_argument('--dim', type = float, default = 128) 17 | parser.add_argument('--epochs', type = int, default = 1000) 18 | parser.add_argument('--valid_epochs', type = int, default = 50) 19 | parser.add_argument('--lamb', type = float, default = 0.01) 20 | parser.add_argument('--data', type = str, default = None) 21 | parser.add_argument('--lr', type = float, default = 2.0) 22 | parser.add_argument('--margin', type = float, default = 1.0) 23 | parser.add_argument('--test', type = int, default = 0) 24 | args = parser.parse_args() 25 | 26 | random.seed(1234) 27 | trainData = TrainData( 28 | path_data = "../data/"+ args.data +"/", 29 | num_neg = args.neg 30 | ) 31 | 32 | model = TransE_type( 33 | num_ent = trainData.num_ent, 34 | num_rel = trainData.num_rel, 35 | dim = args.dim, 36 | p_norm = 2 37 | ) 38 | 39 | model.cuda() 40 | 41 | trainDataLoader = DataLoader(trainData, batch_size = len(trainData) // 100, shuffle = True) 42 | 43 | validData = ValidData( 44 | path_data = "../data/"+ args.data +"/", 45 | test = args.test 46 | ) 47 | 48 | validDataLoader = DataLoader(validData, batch_size = 1, shuffle = False) 49 | 50 | optimizer = torch.optim.SGD(model.parameters(), lr = args.lr) 51 | loss = nn.MarginRankingLoss(margin = args.margin) 52 | 53 | train = tqdm(range(args.epochs)) 54 | for epoch in train: 55 | total_loss = 0.0 56 | for i, (pos, neg) in enumerate(trainDataLoader): 57 | batch = [pos] 58 | for item in neg: 59 | batch.append(item) 60 | batch = torch.cat(batch, 0) 61 | score = model(batch.cuda()) 62 | score_pos = score[:len(pos)] 63 | score_pos = torch.repeat_interleave(score_pos, trainData.num_neg) 64 | score_neg = score[len(pos):] 65 | output = loss(score_pos, score_neg, torch.ones(len(score_pos)).cuda()) 66 | 67 | ent_embeddings = F.normalize(model.ent_embeddings.weight.data, 2, -1) 68 | # Similarity 69 | center = [] 70 | for key in trainData.type2id: 71 | center.append(torch.mean(ent_embeddings[trainData.type2id[key]], axis=0)) 72 | 73 | c = torch.zeros(model.num_ent, model.dim).cuda() 74 | for idx, key in enumerate(trainData.type2id): 75 | c[trainData.type2id[key]] = center[idx] 76 | loss_sim = torch.mean(torch.norm(c - ent_embeddings, model.p_norm, -1)) 77 | 78 | loss_final = output + args.lamb * loss_sim 79 | loss_final.backward() 80 | optimizer.step() 81 | total_loss += loss_final.item() 82 | 83 | train.set_description("Epoch {} | {:.4f}".format(epoch, total_loss)) 84 | 85 | if epoch % args.valid_epochs == 0: 86 | model.eval() 87 | with torch.no_grad(): 88 | mr = 0.0 89 | mrr = 0.0 90 | hit10 = 0 91 | hit3 = 0 92 | hit1 = 0 93 | for triplet in tqdm(validDataLoader): 94 | h, t, r = triplet[0] 95 | 96 | h = h.item() 97 | t = t.item() 98 | r = r.item() 99 | 100 | type_h = validData.id2type[h] 101 | type_t = validData.id2type[t] 102 | 103 | batch_head = triplet[0].repeat(validData.num_ent, 1) 104 | batch_head[:, 0] = torch.arange(validData.num_ent) 105 | batch_tail = triplet[0].repeat(validData.num_ent, 1) 106 | batch_tail[:, 1] = torch.arange(validData.num_ent) 107 | 108 | score_head = model(batch_head.cuda()).cpu().numpy() 109 | 110 | score_target_head = score_head[h] 111 | score_head[validData.filter_head[(t, r)]] = score_target_head - 1 112 | score_head = score_head[validData.type2id[type_h]] 113 | rank = len(np.where(score_head > score_target_head)[0]) + 1 114 | mr += rank 115 | mrr += 1 / rank 116 | if rank < 2: 117 | hit1 += 1 118 | if rank < 4: 119 | hit3 += 1 120 | if rank < 11: 121 | hit10 += 1 122 | 123 | score_tail = model(batch_tail.cuda()).cpu().numpy() 124 | score_target_tail = score_tail[t] 125 | score_tail[validData.filter_tail[(h, r)]] = score_target_tail - 1 126 | score_tail = score_tail[validData.type2id[type_t]] 127 | rank = len(np.where(score_tail > score_target_tail)[0]) + 1 128 | mr += rank 129 | mrr += 1 / rank 130 | if rank < 2: 131 | hit1 += 1 132 | if rank < 4: 133 | hit3 += 1 134 | if rank < 11: 135 | hit10 += 1 136 | print("[Epoch: {}]".format(epoch)) 137 | print("MR:", mr / (2 * len(validData))) 138 | print("MRR:", mrr / (2 * len(validData))) 139 | print("Hit@10:", hit10 / (2 * len(validData))) 140 | print("Hit@3:", hit3 / (2 * len(validData))) 141 | print("Hit@1:", hit1 / (2 * len(validData))) 142 | model.train() --------------------------------------------------------------------------------