File size: 2,044 Bytes
89650c1
 
 
 
 
 
 
 
 
 
 
 
5d47f24
 
556ac6a
 
 
89650c1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c04ccc0
89650c1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c04ccc0
89650c1
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import os
import sys
from transformers import PretrainedConfig, PreTrainedModel
#sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from ultra.models import Ultra
from ultra.datasets import WN18RR, CoDExSmall, FB15k237, FB15k237Inductive
from ultra.eval import test


class UltraConfig(PretrainedConfig):

    model_type = "ultra"

    auto_map = {
        "AutoConfig": "modeling.UltraConfig",
        "AutoModel": "modeling.UltraForKnowledgeGraphReasoning",
    }

    def __init__(
            self,
            relation_model_layers: int = 6,
            relation_model_dim: int = 64,
            entity_model_layers: int = 6, 
            entity_model_dim: int = 64, 
            **kwargs):
        
        self.relation_model_cfg = dict(
            input_dim=relation_model_dim,
            hidden_dims=[relation_model_dim]*relation_model_layers,
            message_func="distmult",
            aggregate_func="sum",
            short_cut=True,
            layer_norm=True
        )

        self.entity_model_cfg = dict(
            input_dim=entity_model_dim,
            hidden_dims=[entity_model_dim]*entity_model_layers,
            message_func="distmult",
            aggregate_func="sum",
            short_cut=True,
            layer_norm=True
        )

        super().__init__(**kwargs)

class UltraForKnowledgeGraphReasoning(PreTrainedModel):

    config_class = UltraConfig

    def __init__(self, config):
        super().__init__(config)

        self.model = Ultra(
            rel_model_cfg=config.relation_model_cfg,
            entity_model_cfg=config.entity_model_cfg,
        )

    def forward(self, data, batch):
        # data: PyG data object 
        # batch shape: (bs, 1+num_negs, 3)
        return self.model.forward(data, batch)
    

if __name__ == "__main__":

    model = UltraForKnowledgeGraphReasoning.from_pretrained("mgalkin/ultra_50g")
    dataset = CoDExSmall(root="./datasets/")
    test(model, mode="test", dataset=dataset, gpus=None)
    # mrr:      0.497697
    # hits@10:  0.685175