| 47 | Link prediction network for graphs / knowledge graphs |
| 48 | """ |
| 49 | def __init__(self, score_function, *embeddings, **kwargs): |
| 50 | super(LinkPredictor, self).__init__() |
| 51 | if isinstance(score_function, types.FunctionType): |
| 52 | self.score_function = score_function |
| 53 | else: |
| 54 | self.score_function = getattr(LinkPredictor, score_function) |
| 55 | self.kwargs = kwargs |
| 56 | self.embeddings = nn.ModuleList() |
| 57 | for embedding in embeddings: |
| 58 | embedding = torch.as_tensor(embedding) |
| 59 | embedding = nn.Embedding.from_pretrained(embedding, freeze=True) |
| 60 | self.embeddings.append(embedding) |
| 61 | |
| 62 | def forward(self, *indexes): |
| 63 | assert len(indexes) == len(self.embeddings) |