dicee.models.flock

Pure PyTorch Flock, compatible with both official jw9730/flock checkpoints.

Architecture adapted from https://github.com/jw9730/flock at UPSTREAM_COMMIT. See docs/flock.md for random-walk and numerical parity details.

MIT License

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

Attributes

UPSTREAM_COMMIT

Classes

`RMSNorm`	Base class for all neural network modules.
`FeedForward`	Base class for all neural network modules.
`BidirectionalGRU`	Base class for all neural network modules.
`FlockBase`	Shared graph lifecycle and entity-scoring interfaces for graph foundation models.
`Flock`	Entity-prediction Flock with official `flock_entity.pth` parameters.
`FlockRelation`	Relation-prediction Flock with official `flock_relation.pth` parameters.

Functions

consensus(x, logits, ids, num_ids)

Stable multihead softmax pooling over all occurrences of each graph ID.

Module Contents

dicee.models.flock.UPSTREAM_COMMIT = 'f35103d25a78bdf4075de5c673a51de4979aa4d7'

class dicee.models.flock.RMSNorm(dim, eps=1e-05)[source]

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

eps = 1e-05

inference_backend = 'auto'

weight

forward(x)[source]

class dicee.models.flock.FeedForward(dim)[source]

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

w1

w2

w3

forward(x)[source]

class dicee.models.flock.BidirectionalGRU(dim, layers)[source]

Bases: torch.nn.Module

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:: training (bool) – Boolean represents whether this module is in training or evaluation mode.

gru_norm

gru

gru_out

ffn_norm

feed_forward

forward(x)[source]

dicee.models.flock.consensus(x, logits, ids, num_ids)[source]: Stable multihead softmax pooling over all occurrences of each graph ID.

class dicee.models.flock.FlockBase(args)[source]

Bases: dicee.models.graph_model.GraphKGE

Shared graph lifecycle and entity-scoring interfaces for graph foundation models.

config_prefix = 'flock'

graph_filename = 'flock_graph.pt'

relation_prediction = False

checkpoint_hint = 'official Flock requires dim=64, walk_len=128, refinements=6, num_layers=1, attention_heads=4...

emb_anon_node: torch.nn.ModuleList

emb_anon_type: torch.nn.ModuleList

emb_restart: torch.nn.ModuleList

emb_neighbor: torch.nn.ModuleList

emb_direction: torch.nn.ModuleList

emb_head_is_query: torch.nn.ModuleList

emb_tail_is_query: torch.nn.ModuleList

emb_node_is_query: torch.nn.ModuleList

emb_type_is_query: torch.nn.ModuleList

from_node: torch.nn.ModuleList

from_type: torch.nn.ModuleList

to_node: torch.nn.ModuleList

to_type: torch.nn.ModuleList

node_logit: torch.nn.ModuleList

type_logit: torch.nn.ModuleList

dim

walk_num

walk_len

refinements

test_samples

seed

prefetch_walks

compact_state

compile_sampler

pack_walks

node_init

type_init

net

head

sample_walks(heads, tails=None, edges=None, generator=None)[source]

Sample official-format records; IDs use the attached graph vocabulary.

Returns (nodes, anonymous nodes, restarts, neighbors, relations, anonymous relations, directions), each shaped [T,B,S,L]. Anonymous names start at 1; no-relation markers are 2*R and L+1 respectively.

score_walks(heads, query, candidates, records)[source]

Score fixed walks for parity/replay; query/candidates use internal IDs.

Entity mode: query=relations, candidates=tails. Relation mode: query=tails, candidates=relations. This scores the supplied records; callers constructing training records must first remove target edges.

class dicee.models.flock.Flock(args)[source]

Bases: FlockBase

Entity-prediction Flock with official flock_entity.pth parameters.

name = 'Flock'

class dicee.models.flock.FlockRelation(args)[source]

Bases: FlockBase, dicee.models.graph_model.RelationGraphKGE

Relation-prediction Flock with official flock_relation.pth parameters.

name = 'FlockRelation'

relation_prediction = True