import logging import os import random import threading from pathlib import Path import yaml from django.conf import settings from api.services.constants import COINS_CONFIG_SUFFIX, COINS_DATASET_META from api.utils import clean_entity_name, clean_relation_name logger = logging.getLogger(__name__) def _safe_load_lightning_checkpoint(cls, ckpt_path): """Load a PyTorch Lightning checkpoint without triggering DDP or deepcopy crashes. Bypasses ``load_from_checkpoint`` entirely: we torch.load the checkpoint to CPU, extract the hyper-parameters, construct the model, load the state_dict, and then move to the target device. This avoids: - DDP __setstate__/__getstate__ needing a process group - save_hyperparameters deepcopy crashing on pickled DDP/datamodule objects - CUDA OOM from loading the full checkpoint (with heavy hparams) onto GPU """ import torch import torch.nn.parallel.distributed as _ddp_mod # Patch DDP so unpickling doesn't need a process group _orig_set = _ddp_mod.DistributedDataParallel.__setstate__ _orig_get = _ddp_mod.DistributedDataParallel.__getstate__ _ddp_mod.DistributedDataParallel.__setstate__ = lambda self, state: self.__dict__.update(state) _ddp_mod.DistributedDataParallel.__getstate__ = lambda self: self.__dict__ try: # Always load to CPU first — the checkpoint contains heavy hparams objects # (dataset_infos with full datamodule + DDP wrappers) that we don't want on GPU. ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=False) finally: _ddp_mod.DistributedDataParallel.__setstate__ = _orig_set _ddp_mod.DistributedDataParallel.__getstate__ = _orig_get hparams = ckpt.get("hyper_parameters", {}) hparams["train_metrics"] = None hparams["sampling_metrics"] = None hparams["visualization_tools"] = None # Construct model without save_hyperparameters (no-op patch on the class) _orig_save = cls.save_hyperparameters cls.save_hyperparameters = lambda self, *a, **kw: None try: model = cls(**hparams) finally: cls.save_hyperparameters = _orig_save # Load weights and move to target device model.load_state_dict(ckpt["state_dict"], strict=False) del ckpt model.to(settings.TORCH_DEVICE) model.eval() return model # Hugging Face Hub model repo holding all checkpoints. The repo mirrors the # on-disk layout under settings.CHECKPOINTS_ROOT (RESEARCH_ROOT by default), so # snapshot_download() drops every file into its final location and the scan # routines below find them unchanged. HF_CHECKPOINTS_REPO = os.environ.get("HF_CHECKPOINTS_REPO", "Bani57/checkpoints") # Per-area checkpoint subdirectories (relative to CHECKPOINTS_ROOT). Used to # detect a fully-populated tree so we can skip the network round-trip on warm # starts. _CHECKPOINT_SUBDIRS = ( Path("COINs-KGGeneration") / "graph_completion" / "checkpoints", Path("COINs-KGGeneration") / "graph_generation" / "checkpoints", Path("MultiProxAn") / "checkpoints", ) # Shared sampler hyperparameters used across all COINs experiments _SAMPLER_HPARS = { "query_structure": ["1p"], "num_negative_samples": 128, "num_neighbours": 10, "random_walk_length": 10, "context_radius": 2, "pagerank_importances": True, "walks_relation_specific": True, } # Per-dataset base config (leiden resolution, loader hyperparams) _DATASET_BASE = { "freebase": { "leiden_resolution": 5.0e-3, "loader_hpars": { "dataset_name": "freebase", "simulated": False, "sample_source": "smore", "sampler_hpars": _SAMPLER_HPARS, }, }, "wordnet": { "leiden_resolution": 0.0, # ExperimentHpars default (wordnet.yml has no leiden_resolution field) "loader_hpars": { "dataset_name": "wordnet", "simulated": False, "sample_source": "smore", "sampler_hpars": _SAMPLER_HPARS, }, }, "nell": { "leiden_resolution": 2.0e-5, "loader_hpars": { "dataset_name": "nell", "simulated": False, "sample_source": "smore", "sampler_hpars": _SAMPLER_HPARS, }, }, } # Training seed per (dataset, algorithm), from PhD Thesis/Code experiment runs. # Verified by matching num_communities in test_log.txt against checkpoint embedding shapes. # Groups sharing the same community structure: # freebase: transe/distmult/complex/rotate (1092 com) | q2b (1030) | kbgat (1025) # wordnet: transe/distmult/complex/rotate (66 com) | q2b (74) | kbgat (88) # nell: transe/distmult/rotate (282 com) | complex (188) | q2b (282, diff sizes) | kbgat (275) _CHECKPOINT_SEEDS = { ("freebase", "transe"): 4089853924, ("freebase", "distmult"): 4089853924, ("freebase", "complex"): 4089853924, ("freebase", "rotate"): 4089853924, ("freebase", "q2b"): 1503136574, ("freebase", "kbgat"): 123456789, ("wordnet", "transe"): 1919180054, ("wordnet", "distmult"): 1919180054, ("wordnet", "complex"): 1919180054, ("wordnet", "rotate"): 1919180054, ("wordnet", "q2b"): 3312854056, ("wordnet", "kbgat"): 123456789, ("nell", "transe"): 3192206669, ("nell", "distmult"): 3192206669, ("nell", "rotate"): 3192206669, ("nell", "complex"): 2409194445, ("nell", "q2b"): 3793326028, ("nell", "kbgat"): 123456789, } # Vanilla seeds per dataset — used for metadata Loaders (entity/relation search, sample triples). VANILLA_SEEDS = { "freebase": 4089853924, "wordnet": 1919180054, "nell": 3192206669, } def get_checkpoint_config(dataset_id, algorithm): """Return the full config for a specific (dataset, algorithm) checkpoint.""" base = _DATASET_BASE[dataset_id] seed = _CHECKPOINT_SEEDS.get((dataset_id, algorithm)) if seed is None: seed = VANILLA_SEEDS[dataset_id] return {"seed": seed, **base} def _adapt_shape_mismatches(ckpt_state_dict, model_state_dict): """Fix ±1 shape mismatches between checkpoint and current model state_dicts. These arise from minor dataset version differences (e.g. NELL gained/lost one node type or one relation between the training run and the current data load) or from update_state_dict adding an extra padding column that wasn't there at training time. Safe to trim because the removed slice corresponded to an absent type/relation that was never looked up at inference. """ import torch as pt adapted = dict(ckpt_state_dict) for key, model_w in model_state_dict.items(): if key not in adapted: continue ckpt_w = adapted[key] if ckpt_w.shape == model_w.shape: continue v = ckpt_w for dim in range(min(len(model_w.shape), len(v.shape))): if v.shape[dim] == model_w.shape[dim] + 1: slc = [slice(None)] * len(v.shape) slc[dim] = slice(None, model_w.shape[dim]) v = v[tuple(slc)] logger.debug("Trimmed %s dim %d: %s -> %s", key, dim, ckpt_w.shape, v.shape) elif model_w.shape[dim] == v.shape[dim] + 1: pad_shape = list(v.shape) pad_shape[dim] = 1 v = pt.cat([v, pt.zeros(pad_shape, dtype=v.dtype)], dim=dim) logger.debug("Padded %s dim %d: %s -> %s", key, dim, ckpt_w.shape, v.shape) adapted[key] = v return adapted def _adapt_mlp_bn_keys(state_dict): """Rename MLP BatchNorm keys from torch_geometric 2.0.x to 2.3.x format. In torch_geometric 2.0.x, MLP stored BN directly as ``norms.N.weight``. In 2.3.x the BN is wrapped in a ModuleList proxy, producing ``norms.N.module.weight``. This affects all MLP BN parameters: weight, bias, running_mean, running_var, num_batches_tracked. Renaming restores the trained BN statistics (running_mean / running_var can differ substantially from the (0, 1) defaults, which is why loading without this fix produced near-zero Q2B scores). """ import re adapted = {} _bn_pattern = re.compile(r"(\.norms\.\d+)\.(weight|bias|running_mean|running_var|num_batches_tracked)$") for key, value in state_dict.items(): new_key = _bn_pattern.sub(r"\1.module.\2", key) adapted[new_key] = value return adapted def _adapt_kbgat_state_dict(ckpt_state_dict, model_state_dict): """Adapt a KBGAT embedder checkpoint from torch_geometric 2.0.x to 2.3.x format. In torch_geometric 2.0.x, the final GATConv layer's out_channels was interpreted as the *total* output width (divided by num_heads internally). In 2.3.x it is per-head, so the weight matrices in the last conv are num_heads-times wider. **Accuracy impact of the repeat strategy**: Both new attention heads receive the full original weight matrix, making them identical. Multi-head diversity is lost — the model degenerates to single-head attention with doubled hidden width. Predictions remain directionally correct (the trained linear transformation is preserved) but may not match the published benchmark numbers exactly. """ import torch as pt adapted = {} for key, value in ckpt_state_dict.items(): if key not in model_state_dict: adapted[key] = value continue expected_shape = model_state_dict[key].shape if value.shape == expected_shape: adapted[key] = value continue # For each dimension where the model expects exactly 2× the checkpoint size, repeat. # (±1 off-by-ones are already handled by _adapt_shape_mismatches before this runs.) v = value for dim in range(len(expected_shape)): if dim < len(v.shape) and expected_shape[dim] == 2 * v.shape[dim]: v = pt.cat([v, v], dim=dim) adapted[key] = v return adapted def _free_heavy_arrays(loader): """Free memory-intensive arrays from a Loader that aren't needed for discovery endpoints.""" loader.node_neighbours = None loader.com_neighbours = None loader.node_adjacency = None loader.com_adjacency = None loader.label_community_edge_freqs = None loader.label_community_edge_freqs_index = None loader.machines = None loader.graph = None loader.node_importances = None loader.neighbour_importances = None loader.out_degrees = None loader.in_degrees = None loader.degrees = None loader.node_degree_type_freqs = None loader.relation_freqs = None class SubgraphInfo: """Holds pre-computed sample subgraphs for a KG anomaly dataset.""" __slots__ = ("subgraphs",) def __init__(self, subgraphs): self.subgraphs = subgraphs # --------------------------------------------------------------------------- # Query sampling — delegates to Query.instantiate() from the research code. # # _STRUCTURE_INFO maps frontend query structure IDs to: # - research code structure string (e.g. "ip" -> "2i1p") # - anchor_map: {query_tree_node_index: frontend_anchor_id} # - variable_map: {query_tree_node_index: frontend_variable_id} # - relation_map: {mapped_tree_edge_index: frontend_relation_id} for "p" edges # # After Query.instantiate() + map_to_tree(), entities live on vertices # and relation strings (like "p5") live on edges indexed by child node index. # --------------------------------------------------------------------------- # (structure, anchor_map, variable_map, relation_map) _STRUCTURE_INFO = { "1p": ("1p", {0: "a"}, {}, {0: "r1"}), "2p": ("2p", {0: "a"}, {1: "v1"}, {0: "r1", 1: "r2"}), "3p": ("3p", {0: "a"}, {1: "v1", 2: "v2"}, {0: "r1", 1: "r2", 2: "r3"}), "2i": ("2i", {0: "a1", 2: "a2"}, {}, {0: "r1", 2: "r2"}), "3i": ("3i", {0: "a1", 2: "a2", 4: "a3"}, {}, {0: "r1", 2: "r2", 4: "r3"}), "ip": ("2i1p", {0: "a1", 2: "a2"}, {4: "v1"}, {0: "r1", 2: "r2", 4: "r3"}), "pi": ("1p2i", {0: "a1", 3: "a2"}, {1: "v1"}, {0: "r1", 1: "r2", 3: "r3"}), } class ModelRegistry: _instance = None def __init__(self): self.coins_checkpoints_available = {} self.graphgen_checkpoints_available = {} self.kg_anomaly_checkpoints_available = {} self.loaders = {} # dataset_id -> lightweight Loader for discovery self.kg_anomaly_subgraphs = {} # dataset_id -> SubgraphInfo self._inference_lock = threading.Lock() self._inference_lock_owner = None # description of who holds the lock self._coins_experiments = {} # (dataset_id, algorithm) -> Experiment self._coins_loaders = {} # (dataset_id, seed, leiden_resolution) -> full Loader self._graphgen_models = {} # (dataset_id, model_type) -> loaded eval-mode model self._kg_anomaly_models = {} # (dataset_id, task) -> loaded eval-mode model def force_release_inference_lock(self): """Emergency release for a stuck inference lock (e.g. client disconnect).""" if self._inference_lock.locked(): self._inference_lock.release() owner = self._inference_lock_owner self._inference_lock_owner = None logger.warning("Inference lock force-released (was held by: %s)", owner) return True return False @classmethod def get(cls): if cls._instance is None: raise RuntimeError("ModelRegistry not initialized. Call initialize() first.") return cls._instance @classmethod def initialize(cls): if cls._instance is not None: return instance = cls() instance._download_checkpoints() instance._scan_checkpoints() instance._load_all_loaders() instance._generate_sample_subgraphs() cls._instance = instance logger.info( "ModelRegistry initialized: coins=%s, multiproxan=%s, kg_anomaly=%s, loaders=%s", instance.is_coins_loaded(), instance.is_graphgen_loaded(), instance.is_kg_anomaly_loaded(), list(instance.loaders.keys()), ) # ---- Checkpoint download ------------------------------------------- def _download_checkpoints(self): """Download checkpoints from Hugging Face Hub if not already present. The HF repo mirrors the on-disk layout under ``CHECKPOINTS_ROOT``, so a single ``snapshot_download`` drops every file into its final location. Idempotent: when all expected subdirs are populated we skip the network round-trip. In production the entrypoint script also pre-warms this download before gunicorn starts, so workers never block on it. """ if self._all_checkpoint_dirs_populated(): logger.info("All checkpoint directories already populated, skipping HF Hub download") return try: from huggingface_hub import snapshot_download except ImportError: logger.warning("huggingface_hub not installed, skipping checkpoint download") return target = Path(settings.CHECKPOINTS_ROOT) target.mkdir(parents=True, exist_ok=True) logger.info("Downloading checkpoints from HF Hub repo %s -> %s", HF_CHECKPOINTS_REPO, target) try: snapshot_download( repo_id=HF_CHECKPOINTS_REPO, repo_type="model", local_dir=str(target), local_dir_use_symlinks=False, max_workers=4, token=os.environ.get("HF_TOKEN"), ) except Exception: logger.exception("Failed to download checkpoints from HF Hub, continuing with local files") def _all_checkpoint_dirs_populated(self): """True if every expected checkpoint subdir contains at least one weight file.""" root = Path(settings.CHECKPOINTS_ROOT) for sub in _CHECKPOINT_SUBDIRS: dest_dir = root / sub if not dest_dir.exists(): return False ckpt_files = list(dest_dir.glob("*.tar")) + list(dest_dir.glob("*.ckpt")) if not ckpt_files: return False return True # ---- Checkpoint scanning ------------------------------------------- def _scan_checkpoints(self): self._scan_coins_checkpoints() self._scan_graphgen_checkpoints() self._scan_kg_anomaly_checkpoints() def _scan_coins_checkpoints(self): ckpt_dir = Path(settings.COINS_COMPLETION_DIR) / "checkpoints" if not ckpt_dir.exists(): logger.warning("COINs checkpoint dir not found: %s", ckpt_dir) return for path in ckpt_dir.glob("*.tar"): parts = path.stem.rsplit("_", 1) if len(parts) == 2: dataset_id, algorithm = parts self.coins_checkpoints_available.setdefault(dataset_id, []).append(algorithm) logger.info("COINs checkpoints: %s", self.coins_checkpoints_available) def _scan_graphgen_checkpoints(self): ckpt_dir = Path(settings.MULTIPROXAN_DIR) / "checkpoints" if not ckpt_dir.exists(): logger.warning("MultiProxAn checkpoint dir not found: %s", ckpt_dir) return for path in ckpt_dir.glob("*.ckpt"): name = path.stem if name.endswith("_c"): dataset_id = name[:-2] self.graphgen_checkpoints_available.setdefault(dataset_id, []).append("continuous") else: self.graphgen_checkpoints_available.setdefault(name, []).append("discrete") logger.info("MultiProxAn checkpoints: %s", self.graphgen_checkpoints_available) def _scan_kg_anomaly_checkpoints(self): ckpt_dir = Path(settings.DIGRESS_KG_DIR) / "checkpoints" if not ckpt_dir.exists(): logger.warning("DiGress KG checkpoint dir not found: %s", ckpt_dir) return for path in ckpt_dir.glob("*.ckpt"): name = path.stem if name.endswith("_correct"): dataset_id = name[:-8] self.kg_anomaly_checkpoints_available.setdefault(dataset_id, []).append("correct") else: self.kg_anomaly_checkpoints_available.setdefault(name, []).append("generate") logger.info("DiGress KG checkpoints: %s", self.kg_anomaly_checkpoints_available) # ---- Loader initialization ----------------------------------------- def _load_all_loaders(self): """Initialize one lightweight Loader per dataset for discovery endpoints. Loads dataset, name maps, train/val/test split, and graph indexes. Heavy arrays (node_neighbours, com_neighbours, adjacency dicts) are freed after startup to save memory. Full Loaders for inference are loaded on demand. """ coins_root = str(Path(settings.COINS_DATA_DIR).parent) original_cwd = os.getcwd() try: os.chdir(coins_root) from graph_completion.graphs.load_graph import Loader, LoaderHpars for dataset_id in _DATASET_BASE: seed = VANILLA_SEEDS[dataset_id] config = get_checkpoint_config(dataset_id, "transe") try: logger.info("Initializing Loader for %s (seed=%d)...", dataset_id, seed) loader = LoaderHpars.from_dict(config["loader_hpars"]).make() leiden_resolution = config["leiden_resolution"] if leiden_resolution is None: dataset_obj = Loader.datasets[dataset_id] dataset_obj.load_from_disk() leiden_resolution = 1.0 / len(dataset_obj.node_data) dataset_obj.unload_from_memory() loader.load_graph( seed=seed, device="cpu", val_size=0.01, test_size=0.02, community_method="leiden", leiden_resolution=leiden_resolution, ) # Rebuild machines from the current num_communities: cached machines_*.npz # files on disk can be stale relative to the current subgraphing output # (different num_communities), and the website always runs num_machines=1. import numpy as np expected_len = loader.num_communities + 2 if len(loader.machines) != expected_len: logger.warning( "Stale machines.npz for %s: len=%d but num_communities+2=%d; rebuilding", dataset_id, len(loader.machines), expected_len, ) loader.machines = np.zeros(expected_len, dtype=int) self.loaders[dataset_id] = loader # Share this loader with _load_coins_experiment so experiments for the # same (dataset, seed, leiden_resolution) reuse it instead of reloading # the graph. Heavy arrays stay populated — they're needed by full # experiments (embedder/sampler/ranker) and by KG anomaly inference. self._coins_loaders[(dataset_id, seed, leiden_resolution)] = loader logger.info( "Loader ready for %s: %d entities, %d relations, %d train triples", dataset_id, loader.num_nodes, loader.num_relations, len(loader.train_edge_data), ) except Exception: logger.exception("Failed to initialize Loader for %s", dataset_id) finally: os.chdir(original_cwd) # ---- Loader accessor helpers --------------------------------------- def get_loader(self, dataset_id): """Return the metadata Loader for a dataset, or None.""" return self.loaders.get(dataset_id) def get_entity_count(self, dataset_id): loader = self.loaders.get(dataset_id) return loader.num_nodes if loader else 0 def get_relation_count(self, dataset_id): loader = self.loaders.get(dataset_id) return loader.num_relations if loader else 0 def get_inverted_name_maps(self, dataset_id): """Return (inv_node_names, inv_node_types, inv_relation_names) Series for a dataset.""" loader = self.loaders.get(dataset_id) if loader is None: return None, None, None return loader.dataset.get_inverted_name_maps() def search_entities(self, dataset_id, query=None, page=1, page_size=50): """Search entities by substring, return paginated (id, name) list and total.""" loader = self.loaders.get(dataset_id) if loader is None: return [], 0 inv_nodes, _, _ = loader.dataset.get_inverted_name_maps() items = [(int(idx), str(name)) for idx, name in inv_nodes.items()] if query: q = query.lower() items = [ (eid, name) for eid, name in items if q in name.lower() or q in clean_entity_name(name, dataset_id).lower() ] total = len(items) start = (max(1, page) - 1) * page_size return items[start:start + page_size], total def search_relations(self, dataset_id, query=None, page=1, page_size=50): """Search relations by substring, return paginated (id, name) list and total.""" loader = self.loaders.get(dataset_id) if loader is None: return [], 0 _, _, inv_relations = loader.dataset.get_inverted_name_maps() items = [(int(idx), str(name)) for idx, name in inv_relations.items()] if query: q = query.lower() items = [ (rid, name) for rid, name in items if q in name.lower() or q in clean_relation_name(name, dataset_id).lower() ] total = len(items) start = (max(1, page) - 1) * page_size return items[start:start + page_size], total def sample_triples(self, dataset_id, count=10, seed=None): """Return random triples with resolved entity/relation names. When ``seed`` is provided, sampling is deterministic — the same ``(dataset_id, count, seed)`` always yields the same triples. When ``seed`` is None, uses the global RNG. """ loader = self.loaders.get(dataset_id) if loader is None: return [] inv_nodes, _, inv_relations = loader.dataset.get_inverted_name_maps() edge_data = loader.train_edge_data count = min(count, len(edge_data)) rng = random.Random(seed) if seed is not None else random indices = rng.sample(range(len(edge_data)), count) result = [] for i in indices: row = edge_data.iloc[i] h, r, t = int(row.s), int(row.r), int(row.t) h_name = str(inv_nodes.get(h, h)) r_name = str(inv_relations.get(r, r)) t_name = str(inv_nodes.get(t, t)) result.append({ "head": {"id": h, "name": h_name, "label": clean_entity_name(h_name, dataset_id)}, "relation": {"id": r, "name": r_name, "label": clean_relation_name(r_name, dataset_id)}, "tail": {"id": t, "name": t_name, "label": clean_entity_name(t_name, dataset_id)}, }) return result # ---- Query sampling --------------------------------------------------- def sample_query(self, dataset_id, query_structure, count=1, seed=None): """Sample structurally valid queries using Query.instantiate() from the research code. Picks random answer entities, walks the training graph backward via adj_t_to_s to produce fully-instantiated query trees, then extracts anchor entities and relation IDs mapped to frontend slot names. """ import numpy as np loader = self.loaders.get(dataset_id) if loader is None: return [] info = _STRUCTURE_INFO.get(query_structure) if info is None: return [] structure_str, anchor_map, variable_map, relation_map = info adj_t_to_s = loader.graph_indexes[2] inv_nodes, _, inv_relations = loader.dataset.get_inverted_name_maps() coins_root = str(Path(settings.COINS_DATA_DIR).parent) original_cwd = os.getcwd() try: os.chdir(coins_root) from graph_completion.graphs.queries import Query, query_edge_r_to_int finally: os.chdir(original_cwd) query = Query(structure_str) query.build_query_tree() answer_candidates = list(adj_t_to_s.keys()) rng = random.Random(seed) if seed is not None else random np_state = np.random.get_state() if seed is not None: np.random.seed(hash(seed) % (2**32)) def ent(eid): name = str(inv_nodes.get(eid, eid)) return {"id": eid, "name": name, "label": clean_entity_name(name, dataset_id)} def rel(rid): name = str(inv_relations.get(rid, rid)) return {"id": rid, "name": name, "label": clean_relation_name(name, dataset_id)} results = [] max_attempts = count * 200 try: for _ in range(max_attempts): if len(results) >= count: break answer = rng.choice(answer_candidates) qi = next( query.instantiate(adj_t_to_s, loader.num_nodes, loader.num_relations, answer, sample=True), None, ) if qi is None: continue qi_mapped = qi.map_to_tree(query.query_tree) anchors = {} for tree_idx, frontend_id in anchor_map.items(): anchors[frontend_id] = ent(int(qi_mapped.vs[tree_idx]["e"])) variables = {} for tree_idx, frontend_id in variable_map.items(): variables[frontend_id] = ent(int(qi_mapped.vs[tree_idx]["e"])) relations = {} for edge_idx, frontend_id in relation_map.items(): rel_id = query_edge_r_to_int(qi_mapped.es[edge_idx]["r"]) relations[frontend_id] = rel(rel_id) target_id = int(qi_mapped.vs[query.query_answer]["e"]) q = {"anchors": anchors, "relations": relations, "target": ent(target_id)} if variables: q["variables"] = variables results.append(q) finally: np.random.set_state(np_state) return results # ---- Sample subgraph generation ------------------------------------ def _generate_sample_subgraphs(self): """Generate sample subgraphs for KG anomaly using the Loader's context subgraph DFS.""" for dataset_id in COINS_DATASET_META: loader = self.loaders.get(dataset_id) if loader is None: continue try: subgraphs = self._build_sample_subgraphs(dataset_id, loader) self.kg_anomaly_subgraphs[dataset_id] = SubgraphInfo(subgraphs) logger.info("Generated %d sample subgraphs for %s", len(subgraphs), dataset_id) except Exception: logger.exception("Failed to generate sample subgraphs for %s", dataset_id) def _build_sample_subgraphs(self, dataset_id, loader, num_subgraphs=40, max_graph_size=10, seed=None): """Build sample subgraphs using the Sampler's DFS-based context subgraph partitioning. When ``seed`` is provided, the DFS iterates node indices in a shuffled order, so different seeds produce different partitions. Without a seed the order is deterministic (original research-code behaviour). """ inv_nodes, _, inv_relations = loader.dataset.get_inverted_name_maps() node_types = loader.dataset.node_data.type.values def entity_label(idx): raw = inv_nodes.get(idx) if raw is None or raw != raw or str(raw).strip() == "": # None / NaN / empty return f"#{idx}" cleaned = clean_entity_name(str(raw), dataset_id) return cleaned if cleaned else f"#{idx}" def relation_label(idx): raw = inv_relations.get(idx) if raw is None or raw != raw or str(raw).strip() == "": return f"rel#{idx}" cleaned = clean_relation_name(str(raw), dataset_id) return cleaned if cleaned else f"rel#{idx}" # Use the Sampler's DFS partitioning to get context subgraphs samples = loader.sampler.get_context_subgraph_samples_dfs( max_graph_size, loader.graph_indexes, loader.num_nodes, max_samples=num_subgraphs * 5, seed=seed, disable_tqdm=True, ) # Randomly pick (row, col) partition pairs so each sample is structurally # distinct from the others. Without shuffling, the DFS returns samples in # nested (k, l) order, which means the first N samples all reuse # partition 0's nodes. Shuffling + a disjoint-partitions guard gives 5 # different subgraphs each call. import random as _random rng = _random.Random(seed) samples = list(samples) rng.shuffle(samples) used_partitions = set() subgraphs = [] for subgraph_row, subgraph_col, nodes_row, nodes_col, edges in samples: if len(subgraphs) >= num_subgraphs: break if subgraph_row in used_partitions or subgraph_col in used_partitions: continue if len(edges) < 3: continue is_bip = (subgraph_row != subgraph_col) if is_bip: # Inpaint mask math assumes balanced halves (n/4, n/2, 3n/4 split). # Only accept bipartite samples where row/col are the same size and # divisible by 4, so the four quadrants are well-defined. if len(nodes_row) != len(nodes_col) or len(nodes_row) < 2: continue if (2 * len(nodes_row)) % 4 != 0: continue sg_nodes = nodes_row + nodes_col row_size = len(nodes_row) else: if len(nodes_row) < 4 or len(nodes_row) % 2 != 0: continue sg_nodes = nodes_row row_size = len(nodes_row) node_idx = {n: i for i, n in enumerate(sg_nodes)} nodes = [] for n in sg_nodes: type_id = int(node_types[n]) if n < len(node_types) else 0 nodes.append({ "entity_id": n, "entity_name": entity_label(n), "type_id": type_id, }) edge_list = [] for h, r, t in edges: if h in node_idx and t in node_idx: edge_list.append({ "source_idx": node_idx[h], "target_idx": node_idx[t], "relation_id": r, "relation_name": relation_label(r), "entity_name_source": entity_label(h), "entity_name_target": entity_label(t), }) subgraphs.append({ "id": f"sample_{len(subgraphs) + 1}", "num_nodes": len(nodes), "num_edges": len(edge_list), "is_bip": is_bip, "row_size": row_size, "nodes": nodes, "edges": edge_list, }) used_partitions.add(subgraph_row) if is_bip: used_partitions.add(subgraph_col) # Free the partitioning data stored on the sampler loader.sampler.context_subgraphs_nodes = None loader.sampler.context_subgraphs_edges = None return subgraphs # ---- COINs experiment loading ------------------------------------------ def _load_coins_experiment(self, dataset_id, algorithm): """Lazily load and cache a fully-prepared Experiment for (dataset_id, algorithm).""" key = (dataset_id, algorithm) if key in self._coins_experiments: return self._coins_experiments[key] config = get_checkpoint_config(dataset_id, algorithm) seed = config["seed"] leiden_resolution = config["leiden_resolution"] coins_root = str(Path(settings.COINS_DATA_DIR).parent) configs_dir = Path(settings.COINS_COMPLETION_DIR) / "configs" suffix = COINS_CONFIG_SUFFIX[algorithm] config_path = configs_dir / f"{dataset_id}{suffix}.yml" with open(config_path, "r", encoding="utf-8") as f: yaml_config = yaml.safe_load(f) # Compute leiden_resolution for datasets where it is None (= 1/num_nodes) if leiden_resolution is None: original_cwd = os.getcwd() try: os.chdir(coins_root) from graph_completion.graphs.load_graph import Loader dataset_obj = Loader.datasets[dataset_id] dataset_obj.load_from_disk() leiden_resolution = 1.0 / len(dataset_obj.node_data) dataset_obj.unload_from_memory() finally: os.chdir(original_cwd) hpars = yaml_config | { "seed": seed, "leiden_resolution": leiden_resolution, "device": settings.TORCH_DEVICE, "train": False, "test": False, "results_dir": str(Path(settings.COINS_COMPLETION_DIR) / "results"), } # YAML's embedding_dim is the trained model's dim. Used below to # expand the TransE init when its dim is smaller (e.g. wordnet kbgat # was trained at 200d but wordnet's transe_model.tar is 100d). target_embedding_dim = int(yaml_config.get("embedder_hpars", {}).get("embedding_dim", 100)) original_cwd = os.getcwd() try: os.chdir(coins_root) from graph_completion.experiments import ExperimentHpars, update_state_dict experiment = ExperimentHpars.from_dict(hpars).make() # Always rebuild machines after load_graph: the website always uses num_machines=1 # (single device, no multi-GPU parallelism), so all communities go to device index 0. # Cached .npz files may have the wrong size if the Leiden partition changed. _orig_load = experiment.loader.load_graph def _patched_load(*args, **kwargs): _orig_load(*args, **kwargs) import numpy as np experiment.loader.machines = np.zeros( experiment.loader.num_communities + 2, dtype=int ) experiment.loader.load_graph = _patched_load # Share Loaders across algorithms with the same (dataset, seed, resolution). # Each load_graph() reloads the full graph from disk and recomputes community # structures; reusing a cached Loader avoids this for e.g. all four # transe/distmult/complex/rotate variants that share the same seed on every dataset. # # share_memory monkey-patch: experiment.prepare() constructs the embedder # and immediately calls embedder.share_memory() to share weights across # multi-process training workers. We run single-process inference so the # call is gratuitous, and on Linux containers with a small /dev/shm # (Docker default 64 MB, free HF Spaces tmpfs similar) it raises a Bus # error mid-prepare. No-op the Module.share_memory class method for the # duration of prepare() and restore it after. # # torch.load monkey-patch: experiment.prepare() loads transe_model.tar to # initialise the embedder's entity_embeddings_initial buffers. The KBGAT # __init__ then assigns weight.data = init, which silently re-shapes the # 200d embedding layer (built from YAML embedding_dim) to the 100d transe # init's shape — and the trained checkpoint's load_state_dict afterwards # blows up on the dim mismatch. Expand the transe init from 100d to # target_embedding_dim by repeating along the embedding axis (same trick # _adapt_kbgat_state_dict uses for GATConv multi-head expansion). import torch as _pt import torch.nn as _nn _orig_share_memory = _nn.Module.share_memory _nn.Module.share_memory = lambda self: self _orig_torch_load = _pt.load def _expand_transe_load(*args, **kwargs): state_dict = _orig_torch_load(*args, **kwargs) if not isinstance(state_dict, dict): return state_dict if not any(k.endswith("entity_embeddings.weight") for k in state_dict): return state_dict # Detect transe init dim and bail if it already matches. sample = next(v for k, v in state_dict.items() if k.endswith("entity_embeddings.weight") and hasattr(v, "shape")) src_dim = int(sample.shape[-1]) if src_dim == target_embedding_dim or src_dim == 0: return state_dict if target_embedding_dim % src_dim != 0: logger.warning( "TransE init dim %d not a divisor of YAML embedding_dim %d; " "leaving init unchanged (load_state_dict may fail).", src_dim, target_embedding_dim, ) return state_dict n_repeats = target_embedding_dim // src_dim expanded = {} for key, value in state_dict.items(): if not hasattr(value, "shape") or value.ndim < 1: expanded[key] = value continue # entity_embeddings(_initial).weight: [num_entities, dim] -> repeat dim # r_embeddings_initial.weight: [num_relations, dim] -> repeat dim # r_embeddings.weight: [dim, num_relations] -> repeat dim 0 if key.endswith(("entity_embeddings.weight", "entity_embeddings_initial.weight", "r_embeddings_initial.weight")) and value.shape[-1] == src_dim: expanded[key] = value.repeat(*([1] * (value.ndim - 1)), n_repeats) elif key.endswith("r_embeddings.weight") and value.shape[0] == src_dim: expanded[key] = value.repeat(n_repeats, *([1] * (value.ndim - 1))) else: expanded[key] = value logger.info("Expanded transe init from %dd to %dd (x%d repeat) for %s/%s", src_dim, target_embedding_dim, n_repeats, dataset_id, algorithm) return expanded _pt.load = _expand_transe_load try: loader_key = (dataset_id, seed, leiden_resolution) if loader_key in self._coins_loaders: cached_loader = self._coins_loaders[loader_key] # Defensive: ensure machines length matches current num_communities. # _load_all_loaders already does this, but any future code path that # populates _coins_loaders directly could skip it. import numpy as np expected_len = cached_loader.num_communities + 2 if len(cached_loader.machines) != expected_len: cached_loader.machines = np.zeros(expected_len, dtype=int) experiment.loader = cached_loader # Temporarily replace load_graph with a no-op: prepare() will find all # required attributes (num_nodes, communities, graph_indexes, …) already set. _orig_load_graph = cached_loader.load_graph cached_loader.load_graph = lambda *args, **kwargs: None try: experiment.prepare() finally: cached_loader.load_graph = _orig_load_graph logger.info("Reused shared Loader for %s seed=%d", dataset_id, seed) else: experiment.prepare() self._coins_loaders[loader_key] = experiment.loader logger.info("Cached new Loader for %s seed=%d", dataset_id, seed) finally: _nn.Module.share_memory = _orig_share_memory _pt.load = _orig_torch_load ckpt_path = (Path(settings.COINS_COMPLETION_DIR) / "checkpoints" / f"{dataset_id}_{algorithm}.tar") import torch as pt ckpt = pt.load(str(ckpt_path), map_location=settings.TORCH_DEVICE) # strict=False: newer torch_geometric MLP adds BatchNorm layers by default; # checkpoints trained without them load cleanly — the BN default params are # identity in eval mode so they don't corrupt outputs. # Rename MLP BN keys: torch_geometric 2.0.x stored them as norms.N.weight; # 2.3.x wraps them as norms.N.module.weight. Applies to all algorithms with MLPs. embedder_sd = _adapt_mlp_bn_keys(ckpt["embedder_state_dict"]) embedder_sd = update_state_dict(embedder_sd, experiment.loader.num_relations) # _adapt_shape_mismatches must run AFTER update_state_dict (which may accidentally # re-pad node_type_embedder due to the ".w" substring match in update_state_dict). embedder_sd = _adapt_shape_mismatches(embedder_sd, experiment.embedder.state_dict()) if algorithm == "kbgat": # torch_geometric 2.0.x → 2.3.x: final GATConv out_channels interpretation # changed from total to per-head, causing 2× width mismatch in the last conv. embedder_sd = _adapt_kbgat_state_dict(embedder_sd, experiment.embedder.state_dict()) experiment.embedder.load_state_dict(embedder_sd, strict=False) experiment.link_ranker.load_state_dict(ckpt["link_ranker_state_dict"], strict=False) experiment.embedder.to(settings.TORCH_DEVICE).eval() experiment.link_ranker.to(settings.TORCH_DEVICE).eval() # Experiment.prepare() only copies embedding_loss_hpars.margin to link_ranker_hpars # for transe/rotate. For Q2B (and any other margin-based algorithm not in that list) # the link_ranker is created with the default margin=1.0 instead of the YAML value. # Patch it directly from the YAML config so inference scores match training. loss_margin = yaml_config.get("embedding_loss_hpars", {}).get("margin") if loss_margin is not None and algorithm not in ("transe", "rotate"): for ranker in (experiment.link_ranker.link_ranker, experiment.link_ranker.community_link_ranker): if hasattr(ranker, "margin"): ranker.margin = float(loss_margin) logger.debug("Patched %s link_ranker margin to %s", algorithm, loss_margin) finally: os.chdir(original_cwd) # Build full-KG adjacency (train + val + test) for inference validation. # graph_indexes[1/2] only cover train_edge_data; check_negative / get_node_cut_cache # must see all KG edges so that valid answers are never falsely rejected. full_edge_data = experiment.loader.dataset.edge_data full_adj_s_to_t, full_adj_t_to_s = {}, {} for s, r, t in full_edge_data[["s", "r", "t"]].values: s, r, t = int(s), int(r), int(t) full_adj_s_to_t.setdefault(s, {}).setdefault(r, []).append(t) full_adj_t_to_s.setdefault(t, {}).setdefault(r, []).append(s) experiment.full_adj_s_to_t = full_adj_s_to_t experiment.full_adj_t_to_s = full_adj_t_to_s logger.info("Built full-KG adjacency for %s/%s (%d edges)", dataset_id, algorithm, len(full_edge_data)) self._coins_experiments[key] = experiment logger.info("COINs experiment ready: dataset=%s algorithm=%s", dataset_id, algorithm) return experiment # ---- Graph generation (MultiProxAn) ----------------------------------- def _load_graphgen_model(self, dataset_id, model_type): key = (dataset_id, model_type) if key in self._graphgen_models: return self._graphgen_models[key] # Defer imports: sys.path is configured at Django startup, not at module import time if model_type == "discrete": from diffusion_model_discrete import DiscreteDenoisingDiffusion as cls else: from diffusion_model import LiftedDenoisingDiffusion as cls suffix = "_c" if model_type == "continuous" else "" ckpt_path = Path(settings.MULTIPROXAN_DIR) / "checkpoints" / f"{dataset_id}{suffix}.ckpt" if not ckpt_path.exists(): from api.exceptions import ModelUnavailable raise ModelUnavailable(f"Checkpoint not found: {ckpt_path.name}") logger.info("Loading MultiProxAn model: dataset=%s model_type=%s", dataset_id, model_type) model = _safe_load_lightning_checkpoint(cls, ckpt_path) self._graphgen_models[key] = model logger.info("MultiProxAn model ready: dataset=%s model_type=%s", dataset_id, model_type) return model def graphgen_generate_stream(self, dataset_id, model_type, sampling_mode, num_nodes, diffusion_steps, chain_frames, multiprox_params): """Return a generator of NDJSON dicts (progress + result). Lock acquisition and model loading happen eagerly so errors surface as normal DRF exceptions. The returned generator releases the lock in its ``finally`` block. """ from api.exceptions import InferenceBusy from api.services.graphgen_inference import ( encode_state_blob, run_multiprox_init, run_standard_generation, ) if not self._inference_lock.acquire(blocking=False): raise InferenceBusy() self._inference_lock_owner = f"graphgen_generate {dataset_id}/{model_type}/{sampling_mode}" try: model = self._load_graphgen_model(dataset_id, model_type) except Exception: self._inference_lock_owner = None self._inference_lock.release() raise def _gen(): try: if sampling_mode == "standard": for event in run_standard_generation( model, num_nodes, diffusion_steps, chain_frames, dataset_id): if event["type"] == "result": event.update({ "dataset_id": dataset_id, "model_type": model_type, "sampling_mode": sampling_mode, }) yield event else: n = multiprox_params["n"] m = multiprox_params["m"] t = multiprox_params["t"] t_prime = multiprox_params["t_prime"] gibbs_chain_freq = multiprox_params["gibbs_chain_freq"] for event in run_multiprox_init( model, num_nodes, n, m, t, t_prime, gibbs_chain_freq, dataset_id): if event["type"] == "result": state = event.pop("state") state["model_type"] = model_type event.update({ "step": 0, "round_complete": False, "done": False, "state": encode_state_blob(state), }) yield event finally: self._inference_lock_owner = None self._inference_lock.release() return _gen() def graphgen_continue_stream(self, state_b64): """Return a generator of NDJSON dicts (progress + result). Blob decoding, lock acquisition, and model loading happen eagerly so errors surface as normal DRF exceptions. The returned generator releases the lock in its ``finally`` block. """ from api.exceptions import InferenceBusy, InvalidRequestError from api.services.graphgen_inference import ( decode_state_blob, encode_state_blob, run_multiprox_step, ) try: state = decode_state_blob(state_b64) except ValueError as exc: raise InvalidRequestError(str(exc)) if not self._inference_lock.acquire(blocking=False): raise InferenceBusy() self._inference_lock_owner = f"graphgen_continue {state['dataset_id']}/{state['model_type']}" try: model = self._load_graphgen_model(state["dataset_id"], state["model_type"]) except Exception: self._inference_lock_owner = None self._inference_lock.release() raise def _gen(): try: for event in run_multiprox_step(model, state, state["dataset_id"]): if event["type"] == "result": updated_state = event.pop("state") event.update({ "step": updated_state["step"], "state": encode_state_blob(updated_state), }) yield event finally: self._inference_lock_owner = None self._inference_lock.release() return _gen() # ---- KG anomaly (DiGress KG) inference -------------------------------- def _load_kg_anomaly_model(self, dataset_id, task): """Load the DiGress KG checkpoint for (dataset_id, task), cached. The KG checkpoint pickles only ``cfg`` via ``save_hyperparameters('cfg')``, so we must reconstruct ``dataset_infos``, ``extra_features`` and ``domain_features`` before constructing the model. Dims are inferred from state_dict shapes; kg_experiment comes from the matching COINs experiment. """ key = (dataset_id, task) if key in self._kg_anomaly_models: return self._kg_anomaly_models[key] import torch import torch.nn.parallel.distributed as _ddp_mod suffix = "_correct" if task == "correct" else "" ckpt_path = Path(settings.DIGRESS_KG_DIR) / "checkpoints" / f"{dataset_id}{suffix}.ckpt" if not ckpt_path.exists(): from api.exceptions import ModelUnavailable raise ModelUnavailable(f"KG anomaly checkpoint not found: {ckpt_path.name}") logger.info("Loading KG anomaly model: dataset=%s task=%s", dataset_id, task) # Load to CPU with DDP patching (same strategy as _safe_load_lightning_checkpoint) _orig_set = _ddp_mod.DistributedDataParallel.__setstate__ _orig_get = _ddp_mod.DistributedDataParallel.__getstate__ _ddp_mod.DistributedDataParallel.__setstate__ = lambda self, state: self.__dict__.update(state) _ddp_mod.DistributedDataParallel.__getstate__ = lambda self: self.__dict__ try: ckpt = torch.load(str(ckpt_path), map_location="cpu", weights_only=False) finally: _ddp_mod.DistributedDataParallel.__setstate__ = _orig_set _ddp_mod.DistributedDataParallel.__getstate__ = _orig_get hparams = ckpt.get("hyper_parameters", {}) cfg = hparams.get("cfg") if isinstance(hparams, dict) else getattr(hparams, "cfg", None) if cfg is None: raise RuntimeError(f"KG anomaly checkpoint {ckpt_path.name} is missing 'cfg' in hyper_parameters") state_dict = ckpt["state_dict"] # Ensure the model's task matches the endpoint task. try: cfg.model.task = task except Exception: pass # OmegaConf struct-mode tolerant: if already set, leave it # Infer dims from state_dict edim_output = state_dict["model.mlp_out_E.2.weight"].shape[0] input_dim_x = state_dict["model.mlp_in_X.0.weight"].shape[1] input_dim_e = state_dict["model.mlp_in_E.0.weight"].shape[1] input_dim_y = state_dict["model.mlp_in_y.0.weight"].shape[1] # Load COINs experiment — needed for kg_experiment and for num_node_types experiment = self._load_coins_experiment(dataset_id, "transe") xdim_output = experiment.loader.num_node_types # Sanity: input_dim_e should equal edim_output (no extra E features for KG) if input_dim_e != edim_output: logger.warning( "Unexpected mlp_in_E dim %d != edim_output %d for %s/%s", input_dim_e, edim_output, dataset_id, task, ) # Build mock dataset_infos from graph_generation.src.diffusion.distributions import DistributionNodes from graph_generation.src.diffusion.extra_features import ( DummyExtraFeatures, ExtraFeatures, ) # max_num_nodes from dataset name (e.g. "freebase_20" -> 20, then *2 per kg_dataset.py) try: base_max = int(cfg.dataset.name.split("_")[-1]) except (AttributeError, ValueError): base_max = 20 max_num_nodes = base_max * 2 # Histogram for DistributionNodes — uniform over possible node counts n_hist = torch.ones(max_num_nodes + 1) n_hist[:2] = 0 # at least 2 nodes nodes_dist = DistributionNodes(n_hist) class _MockDataModule: def __init__(self, kg_experiment, max_num_nodes): self.kg_experiment = kg_experiment self.max_num_nodes = max_num_nodes class _MockDatasetInfos: pass dataset_infos = _MockDatasetInfos() dataset_infos.datamodule = _MockDataModule(experiment, max_num_nodes) dataset_infos.input_dims = {"X": input_dim_x, "E": input_dim_e, "y": input_dim_y} dataset_infos.output_dims = {"X": xdim_output, "E": edim_output, "y": 0} dataset_infos.nodes_dist = nodes_dist dataset_infos.max_n_nodes = max_num_nodes dataset_infos.node_types = torch.ones(xdim_output, dtype=torch.float32) dataset_infos.edge_types = torch.ones(edim_output, dtype=torch.float32) # extra_features per cfg extra_features_type = getattr(cfg.model, "extra_features", None) if cfg.model.type == "discrete" and extra_features_type is not None: extra_features = ExtraFeatures(extra_features_type, dataset_info=dataset_infos) else: extra_features = DummyExtraFeatures() domain_features = DummyExtraFeatures() from diffusion_model_discrete_kg import DiscreteDenoisingDiffusionKG as cls _orig_save = cls.save_hyperparameters cls.save_hyperparameters = lambda self, *a, **kw: None try: model = cls(cfg, dataset_infos, None, None, None, extra_features, domain_features) finally: cls.save_hyperparameters = _orig_save missing, unexpected = model.load_state_dict(state_dict, strict=False) if missing: logger.debug("KG anomaly state_dict missing keys: %d (e.g. %s)", len(missing), missing[:3]) if unexpected: logger.debug("KG anomaly state_dict unexpected keys: %d (e.g. %s)", len(unexpected), unexpected[:3]) del ckpt model.to(settings.TORCH_DEVICE) model.eval() self._kg_anomaly_models[key] = model logger.info("KG anomaly model ready: dataset=%s task=%s", dataset_id, task) return model def kg_anomaly_correct_stream(self, dataset_id, task, sampling_mode, subgraph, diffusion_steps, chain_frames, multiprox_params): """Return a generator of SSE event dicts for /kg-anomaly/correct.""" from api.exceptions import InferenceBusy from api.services.kg_anomaly_inference import ( build_kg_tensors, encode_state_blob, run_multiprox_correction_init, run_standard_correction, ) if not self._inference_lock.acquire(blocking=False): raise InferenceBusy() self._inference_lock_owner = f"kg_anomaly_correct {dataset_id}/{task}/{sampling_mode}" try: model = self._load_kg_anomaly_model(dataset_id, task) loader = self.loaders.get(dataset_id) tensors = build_kg_tensors(subgraph, loader, model) except Exception: self._inference_lock_owner = None self._inference_lock.release() raise def _gen(): try: if sampling_mode == "standard": for event in run_standard_correction( model, tensors, dataset_id, task, loader, diffusion_steps, chain_frames): if event["type"] == "result": event.update({ "dataset_id": dataset_id, "task": task, "sampling_mode": sampling_mode, }) yield event else: n = multiprox_params["n"] m = multiprox_params["m"] t = multiprox_params["t"] t_prime = multiprox_params["t_prime"] gibbs_chain_freq = multiprox_params["gibbs_chain_freq"] for event in run_multiprox_correction_init( model, tensors, dataset_id, task, loader, n, m, t, t_prime, gibbs_chain_freq): if event["type"] == "result": state = event.pop("state") event.update({ "dataset_id": dataset_id, "task": task, "sampling_mode": sampling_mode, "step": 0, "round_complete": False, "done": False, "state": encode_state_blob(state), }) yield event finally: self._inference_lock_owner = None self._inference_lock.release() return _gen() def kg_anomaly_continue_stream(self, state_b64): """Return a generator of SSE event dicts for /kg-anomaly/continue.""" from api.exceptions import InferenceBusy, InvalidRequestError from api.services.kg_anomaly_inference import ( decode_state_blob, encode_state_blob, run_multiprox_correction_step, ) try: state = decode_state_blob(state_b64) except ValueError as exc: raise InvalidRequestError(str(exc)) if not self._inference_lock.acquire(blocking=False): raise InferenceBusy() self._inference_lock_owner = ( f"kg_anomaly_continue {state['dataset_id']}/{state['task']}" ) try: model = self._load_kg_anomaly_model(state["dataset_id"], state["task"]) loader = self.loaders.get(state["dataset_id"]) except Exception: self._inference_lock_owner = None self._inference_lock.release() raise def _gen(): try: for event in run_multiprox_correction_step(model, state, loader): if event["type"] == "result": updated_state = event.pop("state") event.update({ "dataset_id": updated_state["dataset_id"], "task": updated_state["task"], "step": updated_state["step"], "state": encode_state_blob(updated_state), }) yield event finally: self._inference_lock_owner = None self._inference_lock.release() return _gen() # ---- COINs inference --------------------------------------------------- def coins_predict(self, dataset_id, algorithm, query_structure_id, anchors, variables, relations_map, top_k): """Run COINs link prediction / query answering for a single query.""" from api.exceptions import InferenceBusy from api.services.coins_inference import coins_predict_inner if not self._inference_lock.acquire(blocking=False): raise InferenceBusy() self._inference_lock_owner = f"coins_predict {dataset_id}/{algorithm}" try: experiment = self._load_coins_experiment(dataset_id, algorithm) return coins_predict_inner( experiment, dataset_id, algorithm, query_structure_id, anchors, variables, relations_map, top_k, ) finally: self._inference_lock_owner = None self._inference_lock.release() # ---- Status -------------------------------------------------------- def is_coins_loaded(self): return bool(self.coins_checkpoints_available) def is_graphgen_loaded(self): return bool(self.graphgen_checkpoints_available) def is_kg_anomaly_loaded(self): return bool(self.kg_anomaly_checkpoints_available)