Synthetic math knowledge graph dataset. Splits: train, valid, test
| 560 | |
| 561 | |
| 562 | class Math(Dataset): |
| 563 | """ |
| 564 | Synthetic math knowledge graph dataset. |
| 565 | |
| 566 | Splits: |
| 567 | train, valid, test |
| 568 | """ |
| 569 | |
| 570 | NUM_ENTITY = 1000 |
| 571 | NUM_RELATION = 30 |
| 572 | OPERATORS = [ |
| 573 | ("+", lambda x, y: (x + y) % Math.NUM_ENTITY), |
| 574 | ("-", lambda x, y: (x - y) % Math.NUM_ENTITY), |
| 575 | ("*", lambda x, y: (x * y) % Math.NUM_ENTITY), |
| 576 | ("/", lambda x, y: x // y), |
| 577 | ("%", lambda x, y: x % y) |
| 578 | ] |
| 579 | |
| 580 | def __init__(self): |
| 581 | super(Math, self).__init__( |
| 582 | "math", |
| 583 | urls={ |
| 584 | "train": [], |
| 585 | "valid": [], |
| 586 | "test": [] |
| 587 | } |
| 588 | ) |
| 589 | |
| 590 | def train_preprocess(self, save_file): |
| 591 | np.random.seed(1023) |
| 592 | self.generate_math(save_file, num_triplet=20000) |
| 593 | |
| 594 | def valid_preprocess(self, save_file): |
| 595 | np.random.seed(1024) |
| 596 | self.generate_math(save_file, num_triplet=1000) |
| 597 | |
| 598 | def test_preprocess(self, save_file): |
| 599 | np.random.seed(1025) |
| 600 | self.generate_math(save_file, num_triplet=1000) |
| 601 | |
| 602 | def generate_math(self, save_file, num_triplet): |
| 603 | with open(save_file, "w") as fout: |
| 604 | for _ in range(num_triplet): |
| 605 | i = int(np.random.rand() * len(self.OPERATORS)) |
| 606 | op, f = self.OPERATORS[i] |
| 607 | x = int(np.random.rand() * self.NUM_ENTITY) |
| 608 | y = int(np.random.rand() * self.NUM_RELATION) + 1 |
| 609 | fout.write("%d\t%s%d\t%d\n" % (x, op, y, f(x, y))) |
| 610 | |
| 611 | |
| 612 | class FB15k(Dataset): |