probabilistic_kg/contexts.py at main · allthingssecurity/probabilistic_kg · GitHub

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
72
73
74
75
76
77
78
79
80
81
82
83
84
from __future__ import annotations
from collections import defaultdict
from dataclasses import dataclass
from typing import Dict, FrozenSet, Iterable, List, Set, Tuple

from .kg import KnowledgeGraph
from .rule_mining import Rule, UnaryRule, apply_unary


ContextId = int


@dataclass
class Context:
    id: ContextId
    # Rules that comprise this context
    binary_rules: FrozenSet[Rule]
    unary_rules: FrozenSet[UnaryRule]
    # Stats for training
    support_facts: int = 0
    negative_hits: int = 0


def rules_that_entail(
    kg: KnowledgeGraph,
    h: str,
    r: str,
    t: str,
    binary_rules: Iterable[Rule],
    unary_rules: Iterable[UnaryRule],
) -> Tuple[Set[Rule], Set[UnaryRule]]:
    brs: Set[Rule] = set()
    urs: Set[UnaryRule] = set()

    # Binary path rules: r1(x,y) & r2(y,z) -> r(x,z)
    for rule in binary_rules:
        if rule.head != r:
            continue
        # find y s.t. (h, r1, y) and (y, r2, t)
        for (h2, y) in kg.facts.get(rule.r1, set()):
            if h2 != h:
                continue
            if (y, t) in kg.facts.get(rule.r2, set()):
                brs.add(rule)
                break

    # Unary rules: r_body(h,t) -> r_head(h,t) or symmetry
    for rule in unary_rules:
        if rule.head != r:
            continue
        for (h2, t2) in kg.facts.get(rule.body, set()):
            hh, tt = apply_unary(rule, h2, t2)
            if hh == h and tt == t:
                urs.add(rule)
                break

    return brs, urs


def discover_contexts(
    kg: KnowledgeGraph,
    train_triples: Iterable[Tuple[str, str, str]],
    binary_rules: Iterable[Rule],
    unary_rules: Iterable[UnaryRule],
) -> Tuple[Dict[ContextId, Context], Dict[Tuple[str, str, str], ContextId]]:
    # For each training triple, collect rules that entail it; group identical rule-sets as contexts.
    context_index: Dict[Tuple[FrozenSet[Rule], FrozenSet[UnaryRule]], ContextId] = {}
    contexts: Dict[ContextId, Context] = {}
    triple_to_ctx: Dict[Tuple[str, str, str], ContextId] = {}

    next_id = 0
    for (h, r, t) in train_triples:
        brs, urs = rules_that_entail(kg, h, r, t, binary_rules, unary_rules)
        key = (frozenset(brs), frozenset(urs))
        if key not in context_index:
            cid = next_id
            next_id += 1
            context_index[key] = cid
            contexts[cid] = Context(id=cid, binary_rules=key[0], unary_rules=key[1], support_facts=0, negative_hits=0)
        cid = context_index[key]
        contexts[cid].support_facts += 1
        triple_to_ctx[(h, r, t)] = cid

    return contexts, triple_to_ctx