oxirs_core/ai/embeddings/
mod.rs

1//! Knowledge Graph Embeddings for RDF
2//!
3//! This module implements various knowledge graph embedding models including
4//! TransE, DistMult, ComplEx, RotatE, and other state-of-the-art approaches.
5
6use crate::model::Triple;
7use anyhow::Result;
8use serde::{Deserialize, Serialize};
9
10pub mod complex;
11pub mod conve;
12pub mod distmult;
13pub mod evaluation;
14pub mod hype;
15pub mod kgtransformer;
16pub mod neural_tensor_network;
17pub mod rotate;
18pub mod simple;
19pub mod transe;
20pub mod tucker;
21
22pub use complex::ComplEx;
23pub use conve::ConvE;
24pub use distmult::DistMult;
25pub use evaluation::{
26    ConfidenceIntervals, KnowledgeGraphMetrics, LinkPredictionMetrics, StatisticalTestResults,
27    TaskBreakdownMetrics, TrainingMetrics,
28};
29pub use hype::HypE;
30pub use kgtransformer::KGTransformer;
31pub use neural_tensor_network::NeuralTensorNetwork;
32pub use rotate::RotatE;
33pub use simple::SimplE;
34pub use transe::TransE;
35pub use tucker::TuckER;
36
37/// Knowledge graph embedding trait
38#[async_trait::async_trait]
39pub trait KnowledgeGraphEmbedding: Send + Sync {
40    /// Generate embeddings for entities and relations
41    async fn generate_embeddings(&self, triples: &[Triple]) -> Result<Vec<Vec<f32>>>;
42
43    /// Score a triple (head, relation, tail)
44    async fn score_triple(&self, head: &str, relation: &str, tail: &str) -> Result<f32>;
45
46    /// Predict missing links
47    async fn predict_links(
48        &self,
49        entities: &[String],
50        relations: &[String],
51    ) -> Result<Vec<(String, String, String, f32)>>;
52
53    /// Get entity embedding
54    async fn get_entity_embedding(&self, entity: &str) -> Result<Vec<f32>>;
55
56    /// Get relation embedding
57    async fn get_relation_embedding(&self, relation: &str) -> Result<Vec<f32>>;
58
59    /// Train the embedding model
60    async fn train(
61        &mut self,
62        triples: &[Triple],
63        config: &TrainingConfig,
64    ) -> Result<TrainingMetrics>;
65
66    /// Save model to file
67    async fn save(&self, path: &str) -> Result<()>;
68
69    /// Load model from file
70    async fn load(&mut self, path: &str) -> Result<()>;
71}
72
73/// Embedding model configuration
74#[derive(Debug, Clone, Serialize, Deserialize)]
75pub struct EmbeddingConfig {
76    /// Model type
77    pub model_type: EmbeddingModelType,
78
79    /// Embedding dimension
80    pub embedding_dim: usize,
81
82    /// Learning rate
83    pub learning_rate: f32,
84
85    /// L2 regularization weight
86    pub l2_weight: f32,
87
88    /// Negative sampling ratio
89    pub negative_sampling_ratio: f32,
90
91    /// Training batch size
92    pub batch_size: usize,
93
94    /// Maximum training epochs
95    pub max_epochs: usize,
96
97    /// Early stopping patience
98    pub patience: usize,
99
100    /// Validation split
101    pub validation_split: f32,
102
103    /// Enable GPU acceleration
104    pub use_gpu: bool,
105
106    /// Random seed
107    pub seed: u64,
108}
109
110impl Default for EmbeddingConfig {
111    fn default() -> Self {
112        Self {
113            model_type: EmbeddingModelType::TransE,
114            embedding_dim: 100,
115            learning_rate: 0.001,
116            l2_weight: 1e-5,
117            negative_sampling_ratio: 1.0,
118            batch_size: 1024,
119            max_epochs: 1000,
120            patience: 50,
121            validation_split: 0.1,
122            use_gpu: true,
123            seed: 42,
124        }
125    }
126}
127
128/// Embedding model types
129#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
130pub enum EmbeddingModelType {
131    /// Translation-based model (Bordes et al., 2013)
132    TransE,
133
134    /// Bilinear model (Yang et al., 2014)
135    DistMult,
136
137    /// Complex embeddings (Trouillon et al., 2016)
138    ComplEx,
139
140    /// Rotation-based model (Sun et al., 2019)
141    RotatE,
142
143    /// Hyperbolic embeddings (Balazevic et al., 2019)
144    HypE,
145
146    /// Tucker decomposition (Balazevic et al., 2019)
147    TuckER,
148
149    /// Convolutional model (Dettmers et al., 2018)
150    ConvE,
151
152    /// Transformer-based model
153    KGTransformer,
154
155    /// Neural tensor network (Socher et al., 2013)
156    NeuralTensorNetwork,
157
158    /// SimplE (Kazemi & Poole, 2018)
159    SimplE,
160}
161
162/// Create embedding model based on configuration
163pub fn create_embedding_model(
164    config: EmbeddingConfig,
165) -> anyhow::Result<std::sync::Arc<dyn KnowledgeGraphEmbedding>> {
166    match config.model_type {
167        EmbeddingModelType::TransE => Ok(std::sync::Arc::new(TransE::new(config))),
168        EmbeddingModelType::DistMult => Ok(std::sync::Arc::new(DistMult::new(config))),
169        EmbeddingModelType::ComplEx => Ok(std::sync::Arc::new(ComplEx::new(config))),
170        EmbeddingModelType::RotatE => Ok(std::sync::Arc::new(RotatE::new(config))),
171        EmbeddingModelType::HypE => Ok(std::sync::Arc::new(HypE::new(config))),
172        EmbeddingModelType::TuckER => Ok(std::sync::Arc::new(TuckER::new(config))),
173        EmbeddingModelType::ConvE => Ok(std::sync::Arc::new(ConvE::new(config))),
174        EmbeddingModelType::KGTransformer => Ok(std::sync::Arc::new(KGTransformer::new(config))),
175        EmbeddingModelType::NeuralTensorNetwork => {
176            Ok(std::sync::Arc::new(NeuralTensorNetwork::new(config)))
177        }
178        EmbeddingModelType::SimplE => Ok(std::sync::Arc::new(SimplE::new(config))),
179    }
180}
181
182/// Training configuration
183#[derive(Debug, Clone, Serialize, Deserialize)]
184pub struct TrainingConfig {
185    /// Batch size
186    pub batch_size: usize,
187
188    /// Learning rate
189    pub learning_rate: f32,
190
191    /// Maximum epochs
192    pub max_epochs: usize,
193
194    /// Validation split
195    pub validation_split: f32,
196
197    /// Early stopping patience
198    pub patience: usize,
199}
200
201impl Default for TrainingConfig {
202    fn default() -> Self {
203        Self {
204            batch_size: 1024,
205            learning_rate: 0.001,
206            max_epochs: 1000,
207            validation_split: 0.1,
208            patience: 50,
209        }
210    }
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216
217    /// Verify that every EmbeddingModelType variant can be constructed without error.
218    #[test]
219    fn test_create_embedding_model_all_variants() {
220        let variants = [
221            EmbeddingModelType::TransE,
222            EmbeddingModelType::DistMult,
223            EmbeddingModelType::ComplEx,
224            EmbeddingModelType::RotatE,
225            EmbeddingModelType::HypE,
226            EmbeddingModelType::TuckER,
227            EmbeddingModelType::ConvE,
228            EmbeddingModelType::KGTransformer,
229            EmbeddingModelType::NeuralTensorNetwork,
230            EmbeddingModelType::SimplE,
231        ];
232
233        for variant in &variants {
234            let config = EmbeddingConfig {
235                model_type: variant.clone(),
236                embedding_dim: 8,
237                ..Default::default()
238            };
239            let result = create_embedding_model(config);
240            assert!(
241                result.is_ok(),
242                "create_embedding_model should succeed for {:?}",
243                variant
244            );
245        }
246    }
247}
oxirs_core/ai/embeddings/mod.rs

oxirs_core/ai/embeddings/
mod.rs
