qualia_core_db/solvers/learning/kg_embedding/
mod.rs

1//! Knowledge-graph embedding (TransE / DistMult / ComplEx / RotatE) — score a
2//! triple `(head, relation, tail)` for plausibility, and rank candidate entities
3//! for **link prediction** over the semantic graph.
4//!
5//! ## Affordability gating (PROJECT RULE — the honest-scope test)
6//!
7//! KG embedding has two halves with wildly different cost:
8//!
9//! * **Scoring / ranking** ([`score`], [`predict`]) — a few dot products per triple.
10//!   Trivially cheap, always present, runs on any device. This is the path a *user*
11//!   exercises: given an already-trained [`EmbeddingTable`], score and rank.
12//! * **Training** ([`train`]) — gradient descent over many epochs and negatives. This
13//!   is the **heavy, run-once** pass: it is structured as an artifact producer that
14//!   runs on capable hardware and is then *distributed* (the trained table), never on
15//!   a user's critical path. It is dispatch-ready (§13): the per-triple score/gradient
16//!   batch is kernel-class `DenseLinear`, with the CPU reference here always present.
17//!
18//! So nothing here forces a user into food-vs-compute: they consume a table; they do
19//! not have to train one.
20//!
21//! ## Honesty
22//!
23//! Every public entry fails closed ([`KgEmbeddingError`]) on a dimension/index
24//! mismatch rather than returning a fabricated score. A score is only ever produced
25//! from real embedding arithmetic.
26
27pub mod predict;
28pub mod score;
29pub mod train;
30
31pub use predict::{hits_at_k, mean_rank, mean_reciprocal_rank, rank_tail, RankFilter};
32pub use score::ScoreModel;
33pub use train::{train, TrainConfig};
34
35/// Fail-closed errors for the embedding library.
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub enum KgEmbeddingError {
38    /// A vector length did not match the model's expected entity/relation dim.
39    InvalidDimension,
40    /// An entity or relation index was out of range for the table.
41    IndexOutOfRange,
42    /// Not enough triples / rank to fit (e.g. empty training set).
43    InsufficientData,
44    /// A configuration was inconsistent (e.g. rank 0, zero epochs with no table).
45    InvalidParameters,
46}
47
48impl core::fmt::Display for KgEmbeddingError {
49    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
50        match self {
51            KgEmbeddingError::InvalidDimension => write!(f, "embedding vector length mismatch"),
52            KgEmbeddingError::IndexOutOfRange => write!(f, "entity/relation index out of range"),
53            KgEmbeddingError::InsufficientData => write!(f, "insufficient data to fit"),
54            KgEmbeddingError::InvalidParameters => write!(f, "invalid embedding configuration"),
55        }
56    }
57}
58impl std::error::Error for KgEmbeddingError {}
59
60/// A trained (or freshly-initialised) embedding table: one vector per entity and one
61/// per relation. The storage length per entity/relation is fixed by the model and the
62/// rank `k` (see [`ScoreModel::dims`]).
63#[derive(Debug, Clone, PartialEq)]
64pub struct EmbeddingTable {
65    pub model: ScoreModel,
66    /// Rank (the conceptual embedding dimension). Storage may be `k` or `2k` per
67    /// vector depending on the model.
68    pub rank: usize,
69    pub ent_dim: usize,
70    pub rel_dim: usize,
71    pub n_entities: usize,
72    pub n_relations: usize,
73    /// `n_entities * ent_dim`, row-major.
74    pub entities: Vec<f64>,
75    /// `n_relations * rel_dim`, row-major.
76    pub relations: Vec<f64>,
77}
78
79impl EmbeddingTable {
80    /// Allocate a zeroed table sized for `model` at rank `k`.
81    pub fn zeros(
82        model: ScoreModel,
83        k: usize,
84        n_entities: usize,
85        n_relations: usize,
86    ) -> Result<Self, KgEmbeddingError> {
87        if k == 0 || n_entities == 0 || n_relations == 0 {
88            return Err(KgEmbeddingError::InvalidParameters);
89        }
90        let (ent_dim, rel_dim) = model.dims(k);
91        Ok(Self {
92            model,
93            rank: k,
94            ent_dim,
95            rel_dim,
96            n_entities,
97            n_relations,
98            entities: vec![0.0; n_entities * ent_dim],
99            relations: vec![0.0; n_relations * rel_dim],
100        })
101    }
102
103    #[inline]
104    pub fn entity(&self, i: usize) -> Result<&[f64], KgEmbeddingError> {
105        if i >= self.n_entities {
106            return Err(KgEmbeddingError::IndexOutOfRange);
107        }
108        Ok(&self.entities[i * self.ent_dim..(i + 1) * self.ent_dim])
109    }
110
111    #[inline]
112    pub fn relation(&self, i: usize) -> Result<&[f64], KgEmbeddingError> {
113        if i >= self.n_relations {
114            return Err(KgEmbeddingError::IndexOutOfRange);
115        }
116        Ok(&self.relations[i * self.rel_dim..(i + 1) * self.rel_dim])
117    }
118
119    #[inline]
120    pub fn entity_mut(&mut self, i: usize) -> &mut [f64] {
121        let d = self.ent_dim;
122        &mut self.entities[i * d..(i + 1) * d]
123    }
124
125    #[inline]
126    pub fn relation_mut(&mut self, i: usize) -> &mut [f64] {
127        let d = self.rel_dim;
128        &mut self.relations[i * d..(i + 1) * d]
129    }
130
131    /// Plausibility score for triple `(h, r, t)` (entity/relation indices). Higher =
132    /// more plausible. Fails closed on out-of-range indices.
133    pub fn score(&self, h: usize, r: usize, t: usize) -> Result<f64, KgEmbeddingError> {
134        let hv = self.entity(h)?;
135        let rv = self.relation(r)?;
136        let tv = self.entity(t)?;
137        self.model.score(hv, rv, tv, self.rank)
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn table_dims_and_accessors() {
147        let t = EmbeddingTable::zeros(ScoreModel::ComplEx, 4, 3, 2).unwrap();
148        assert_eq!(t.ent_dim, 8); // 2k
149        assert_eq!(t.rel_dim, 8);
150        assert_eq!(t.entities.len(), 24);
151        assert!(t.entity(2).is_ok());
152        assert_eq!(t.entity(3).unwrap_err(), KgEmbeddingError::IndexOutOfRange);
153    }
154
155    #[test]
156    fn rotate_table_relation_is_angles_only() {
157        let t = EmbeddingTable::zeros(ScoreModel::RotatE, 5, 2, 2).unwrap();
158        assert_eq!(t.ent_dim, 10); // 2k
159        assert_eq!(t.rel_dim, 5); // k angles
160    }
161
162    #[test]
163    fn zeros_fails_closed_on_degenerate() {
164        assert_eq!(
165            EmbeddingTable::zeros(ScoreModel::TransE { p: 2 }, 0, 1, 1).unwrap_err(),
166            KgEmbeddingError::InvalidParameters
167        );
168    }
169}
qualia_core_db/solvers/learning/kg_embedding/mod.rs

qualia_core_db/solvers/learning/kg_embedding/
mod.rs
