OxiRS Embed - Knowledge Graph Embeddings

Status: v0.3.2 - Released 2026-07-12

✨ Production Release: Production-ready with API stability guarantees. Semantic versioning enforced.

Generate vector embeddings for RDF knowledge graphs enabling semantic similarity search, entity linking, and neural-symbolic AI integration.

Features

Embedding Models

• TransE - Translational distance models for knowledge graphs
• DistMult - Bilinear diagonal models for symmetric relations
• ComplEx - Complex-valued embeddings for asymmetric relations
• RotatE - Rotation-based models in complex space
• HolE - Holographic embeddings using circular correlation (NEW in v0.1.0)
• ConvE - Convolutional 2D neural network embeddings (NEW in v0.1.0)
• TuckER - Tucker decomposition for multi-relational learning
• QuatE - Quaternion embeddings for complex patterns

Advanced Features (NEW in v0.1.0)

• Link Prediction - Predict missing triples (head/tail/relation)
  • Filtered ranking to remove known triples
  • Batch prediction for efficiency
  • Evaluation metrics (MRR, Hits@K, Mean Rank)
• Entity Clustering - Group similar entities
  • K-Means with K-Means++ initialization
  • Hierarchical (agglomerative) clustering
  • DBSCAN (density-based) clustering
  • Spectral clustering
  • Quality metrics (silhouette score, inertia)
• Community Detection - Find communities in knowledge graphs
  • Louvain modularity optimization
  • Label propagation
  • Girvan-Newman edge betweenness
  • Embedding-based detection
• Vector Search - High-performance semantic search (NEW in 0.1.0)
  • Exact search with multiple distance metrics
  • Cosine similarity, Euclidean, dot product, Manhattan
  • Batch search for multiple queries
  • Radius-based filtering
  • Parallel processing support
• Visualization - t-SNE, PCA, UMAP, Random Projection
  • 2D and 3D dimensionality reduction
  • Export to CSV/JSON formats
  • Cluster-aware visualizations
• Interpretability - Model understanding tools
  • Similarity analysis and nearest neighbors
  • Feature importance analysis
  • Counterfactual explanations
  • Embedding space diagnostics
• Mixed Precision Training - FP16/FP32 for faster training
• Model Quantization - Int8/Int4/Binary compression (3-4x size reduction)

Knowledge Graph Embedding

• Entity Embeddings - Generate embeddings for RDF entities
• Relation Embeddings - Embed predicates and relationships
• Graph Embeddings - Whole-graph vector representations
• Contextual Embeddings - Use graph context for better embeddings

Applications

• Semantic Search - Find similar entities by meaning
• Entity Linking - Link mentions to knowledge graph entities
• Relation Prediction - Predict missing relationships
• Clustering - Group similar entities
• Knowledge Graph Completion - Fill missing facts in KGs
• Anomaly Detection - Detect unusual patterns in graphs

Installation

Add to your Cargo.toml:

[dependencies]
oxirs-embed = "0.3.2"

# Enable optional feature groups (advanced models, GPU acceleration, API server, ...)
oxirs-embed = { version = "0.3.2", features = ["advanced-models", "gpu"] }

Quick Start

Basic Entity Embedding

Free-text embedding uses the domain-specific transformer models in biomedical_embeddings (SciBERT, CodeBERT, BioBERT, LegalBERT, FinBERT, ClinicalBERT, ChemBERT):

use oxirs_embed::biomedical_embeddings::{SpecializedTextEmbedding, SpecializedTextModel};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Load a domain-specific text embedding model (SciBERT here)
    let mut model = SpecializedTextEmbedding::new(SpecializedTextEmbedding::scibert_config());

    // Generate an embedding for a piece of text
    let text = "Machine learning researcher specializing in NLP";
    let embedding = model.encode_text(text).await?;

    println!("Embedding dimension: {}", embedding.len());
    println!("Model: {:?}", SpecializedTextModel::SciBERT.model_name());

    Ok(())
}

Knowledge Graph Embedding

Structural knowledge graph embedding trains one of the EmbeddingModel implementations (TransE, DistMult, ComplEx, RotatE, HolE, ConvE, TuckER, QuatE, ...) directly on RDF triples:

use oxirs_embed::{EmbeddingModel, ModelConfig, NamedNode, TransE, Triple};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Create a knowledge graph embedding model
    let config = ModelConfig::default().with_dimensions(128);
    let mut model = TransE::new(config);

    // Add knowledge triples
    model.add_triple(Triple::new(
        NamedNode::new("http://example.org/Person/Alice")?,
        NamedNode::new("http://example.org/knows")?,
        NamedNode::new("http://example.org/Person/Bob")?,
    ))?;

    // Train the model
    let stats = model.train(Some(100)).await?;
    println!("Training completed: {stats:?}");

    // Get the embedding for a specific entity
    let embedding = model.get_entity_embedding("http://example.org/Person/Alice")?;
    println!("Embedding for Alice: {} dimensions", embedding.dimensions);

    Ok(())
}

TransE also implements Default (delegating to ModelConfig::default(), i.e. 100 dimensions, L2 distance, margin 1.0), so TransE::default() works anywhere a plain M: EmbeddingModel + Default bound is needed — for example in generic benchmark/test harnesses. Prefer TransE::new(config) when you need a specific dimension or learning rate.

Semantic Similarity

use oxirs_embed::embedding_aggregator::cosine_similarity;

// Compare two entity embeddings (both `&[f32]`, e.g. from `Vector::values`)
let score = cosine_similarity(&alice_embedding.values, &bob_embedding.values);
println!("Alice <-> Bob similarity: {score:.3}");

For ranked nearest-neighbor search over many entities at once, use the Vector Search index below rather than a single pairwise call.

New Models (v0.1.0)

HolE (Holographic Embeddings)

HolE uses circular correlation to model entity and relation interactions. Effective for capturing symmetric and asymmetric patterns.

use oxirs_embed::{
    models::hole::{HoLE, HoLEConfig},
    EmbeddingModel, ModelConfig, NamedNode, Triple,
};

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    // Configure HolE model
    let config = HoLEConfig {
        base: ModelConfig {
            dimensions: 100,
            learning_rate: 0.01,
            max_epochs: 100,
            ..Default::default()
        },
        regularization: 0.0001,
        margin: 1.0,
        num_negatives: 5,
        use_sigmoid: true,
    };

    let mut model = HoLE::new(config);

    // Add triples
    model.add_triple(Triple::new(
        NamedNode::new("paris")?,
        NamedNode::new("capital_of")?,
        NamedNode::new("france")?,
    ))?;

    // Train
    let stats = model.train(Some(100)).await?;

    // Score triple
    let score = model.score_triple("paris", "capital_of", "france")?;
    println!("Score: {:.4}", score);

    Ok(())
}

ConvE (Convolutional Embeddings)

ConvE uses 2D CNNs for expressive knowledge graph embeddings. Parameter-efficient with shared convolutional filters.

use oxirs_embed::models::conve::{ConvE, ConvEConfig};

let config = ConvEConfig {
    base: ModelConfig {
        dimensions: 200,
        learning_rate: 0.001,
        max_epochs: 100,
        ..Default::default()
    },
    reshape_width: 20,  // 200 / 20 = 10 height
    num_filters: 32,
    kernel_size: 3,
    dropout_rate: 0.3,
    ..Default::default()
};

let mut model = ConvE::new(config);

// Add triples and train as before
model.add_triple(triple)?;
model.train(Some(100)).await?;

Link Prediction

Predict missing entities or relations in knowledge graphs.

use oxirs_embed::link_prediction::{LinkPredictionConfig, LinkPredictor};

// Create predictor
let pred_config = LinkPredictionConfig {
    top_k: 5,
    filter_known_triples: true,
    min_confidence: 0.0,
    parallel: true,
    batch_size: 100,
};

let predictor = LinkPredictor::new(pred_config, model);

// Predict tail entity (object prediction)
let candidates = vec!["bob".to_string(), "charlie".to_string()];
let predictions = predictor.predict_tail("alice", "knows", &candidates)?;

for pred in predictions {
    println!("{} (score: {:.4}, rank: {})", pred.predicted_id, pred.score, pred.rank);
}

// Predict head entity (subject prediction)
let predictions = predictor.predict_head("knows", "bob", &candidates)?;

// Predict relation
let relations = vec!["knows".to_string(), "friend_of".to_string()];
let predictions = predictor.predict_relation("alice", "bob", &relations)?;

Entity Clustering

Group similar entities based on learned embeddings.

use oxirs_embed::clustering::{ClusteringAlgorithm, ClusteringConfig, EntityClustering};
use std::collections::HashMap;

// Extract embeddings
let mut embeddings = HashMap::new();
for entity in model.get_entities() {
    if let Ok(emb) = model.get_entity_embedding(&entity) {
        let array = scirs2_core::ndarray_ext::Array1::from_vec(emb.values);
        embeddings.insert(entity, array);
    }
}

// K-Means clustering
let config = ClusteringConfig {
    algorithm: ClusteringAlgorithm::KMeans,
    num_clusters: 5,
    max_iterations: 100,
    ..Default::default()
};

let mut clustering = EntityClustering::new(config);
let result = clustering.cluster(&embeddings)?;

println!("Silhouette score: {:.3}", result.silhouette_score);
println!("Cluster assignments:");
for (entity, cluster_id) in result.assignments {
    println!("  {} -> Cluster {}", entity, cluster_id);
}

Community Detection

Find communities in knowledge graphs using graph structure and embeddings.

use oxirs_embed::community_detection::{CommunityAlgorithm, CommunityConfig, CommunityDetector};

let config = CommunityConfig {
    algorithm: CommunityAlgorithm::Louvain,
    min_community_size: 2,
    resolution: 1.0,
    ..Default::default()
};

let mut detector = CommunityDetector::new(config);
let result = detector.detect_from_triples(&triples)?;

println!("Modularity: {:.3}", result.modularity);
println!("Found {} communities", result.communities.len());

Vector Search

High-performance semantic search for knowledge graph embeddings.

use oxirs_embed::vector_search::{VectorSearchIndex, SearchConfig, DistanceMetric};

// Build search index
let config = SearchConfig {
    metric: DistanceMetric::Cosine,
    parallel: true,
    normalize: true,
    ..Default::default()
};

let mut index = VectorSearchIndex::new(config);
index.build(&embeddings)?;

// Search for similar entities
let query_embedding = embeddings["iphone"].to_vec();
let results = index.search(&query_embedding, 10)?;

for result in results {
    println!("{}: similarity = {:.3}", result.entity_id, result.score);
}

// Batch search
let queries = vec![query1, query2, query3];
let batch_results = index.batch_search(&queries, 10)?;

// Radius search (find all within distance)
let radius_results = index.radius_search(&query_embedding, 0.5)?;

Visualization

Visualize embeddings in 2D/3D using dimensionality reduction.

use oxirs_embed::visualization::{EmbeddingVisualizer, ReductionMethod, VisualizationConfig};

// PCA visualization
let config = VisualizationConfig {
    method: ReductionMethod::PCA,
    target_dims: 2,
    ..Default::default()
};

let mut visualizer = EmbeddingVisualizer::new(config);
let result = visualizer.visualize(&embeddings)?;

// t-SNE visualization (better for discovering clusters)
let tsne_config = VisualizationConfig {
    method: ReductionMethod::TSNE,
    target_dims: 2,
    tsne_perplexity: 30.0,
    max_iterations: 1000,
    ..Default::default()
};

let mut tsne_viz = EmbeddingVisualizer::new(tsne_config);
let tsne_result = tsne_viz.visualize(&embeddings)?;

// Export to CSV for plotting
for (entity, coords) in &tsne_result.coordinates {
    println!("{},{},{}", entity, coords[0], coords[1]);
}

Interpretability

Understand why models make certain predictions.

use oxirs_embed::interpretability::{InterpretabilityAnalyzer, InterpretabilityConfig, InterpretationMethod};

// Similarity analysis
let config = InterpretabilityConfig {
    method: InterpretationMethod::SimilarityAnalysis,
    top_k: 10,
    ..Default::default()
};

let analyzer = InterpretabilityAnalyzer::new(config);
let analysis = analyzer.similarity_analysis("alice", &embeddings)?;

println!("Most similar to 'alice':");
for (entity, score) in &analysis.similar_entities {
    println!("  {}: {:.3}", entity, score);
}

// Feature importance
let importance_config = InterpretabilityConfig {
    method: InterpretationMethod::FeatureImportance,
    top_k: 10,
    ..Default::default()
};

let imp_analyzer = InterpretabilityAnalyzer::new(importance_config);
let importance = imp_analyzer.feature_importance("alice", &embeddings)?;

// Counterfactual explanations
let counterfactual = analyzer.counterfactual_explanation("alice", "bob", &embeddings)?;
println!("To be like Bob, Alice would need to change {} dimensions",
    counterfactual.required_changes.len());

Supported Embedding Providers

oxirs-embed ships two families of models rather than a single pluggable "provider" facade: structural knowledge-graph embedders (trained locally on your triples) and specialized domain text-embedding models (also local — no network calls).

Structural Knowledge Graph Models

Selected via Cargo features (basic-models is the default):

Feature	Models
basic-models (default)	TransE, ComplEx, DistMult, HoLE
advanced-models	RotatE, ConvE, TuckER, QuatE

Each is a plain struct implementing the EmbeddingModel trait — see Knowledge Graph Embedding above.

Specialized Domain Text Models

oxirs_embed::biomedical_embeddings::SpecializedTextModel covers seven pretrained architectures, each with a matching SpecializedTextEmbedding::{name}_config() helper (scibert_config(), codebert_config(), biobert_config(), ...):

• SciBERT - Scientific literature (768 dimensions)
• CodeBERT - Code and programming languages
• BioBERT - Biomedical literature
• LegalBERT - Legal documents
• FinBERT - Financial texts
• ClinicalBERT - Clinical notes
• ChemBERT - Chemical compounds

use oxirs_embed::biomedical_embeddings::{SpecializedTextConfig, SpecializedTextEmbedding, SpecializedTextModel};

let config = SpecializedTextConfig {
    model_type: SpecializedTextModel::BioBERT,
    ..Default::default()
};
let mut model = SpecializedTextEmbedding::new(config);
let embedding = model.encode_text("BRCA1 is associated with breast cancer risk").await?;

Advanced Features

Batch Processing

The EmbeddingModel::encode trait method already accepts multiple texts per call — there is no separate batch-sized entry point:

use oxirs_embed::EmbeddingModel;

let texts: Vec<String> = vec![
    "Machine learning".to_string(),
    "Natural language processing".to_string(),
    "Computer vision".to_string(),
];

// One call encodes the whole batch: Vec<Vec<f32>>, one embedding per input text
let embeddings = model.encode(&texts).await?;

Contextual Embeddings

ContextualEmbeddingModel adapts embeddings to query/user/task/temporal context. It is still evolving (its crate-root re-export is currently disabled), so import it via its module path:

use oxirs_embed::contextual::{ContextualConfig, ContextualEmbeddingModel, EmbeddingContext};
use oxirs_embed::Triple;

# async fn example(triples: &[Triple]) -> anyhow::Result<()> {
let mut model = ContextualEmbeddingModel::new(ContextualConfig::default())?;
let context = EmbeddingContext::default();

let embeddings = model.embed_with_context(triples, &context).await?;
println!("Generated {} contextual embeddings", embeddings.len());
# Ok(())
# }

Entity Linking

use oxirs_embed::entity_linking::{EntityLinker, EntityLinkerConfig};
use scirs2_core::ndarray_ext::Array1;
use std::collections::HashMap;

let linker = EntityLinker::new(EntityLinkerConfig::default(), entity_embeddings)?;

// Link a (pre-embedded) mention to knowledge graph entities — the mention text
// "machine learning expert from Stanford" must already be embedded into the
// same vector space as `entity_embeddings` (e.g. via a text embedding model).
let mention_embedding: Array1<f32> = embed_mention("machine learning expert from Stanford");
let candidates = linker.link_entity(&mention_embedding, None)?;

for result in candidates {
    println!("{}: confidence={:.3}", result.entity_id, result.confidence);
}

Relation Prediction

use oxirs_embed::entity_linking::{RelationPredictor, RelationPredictorConfig};

let predictor = RelationPredictor::new(
    RelationPredictorConfig::default(),
    relation_embeddings,
    entity_embeddings,
);

// Predict relations between two entities
let predictions = predictor.predict_relations(
    "http://example.org/Alice",
    "http://example.org/Bob",
)?;

for pred in predictions {
    println!("Predicted relation: {} (score={:.3})", pred.relation, pred.score);
}

Integration with OxiRS

With oxirs-vec (Vector Search)

use oxirs_embed::{EmbeddingModel, TransE};
use oxirs_vec::{Vector as VecVector, VectorStore, VectorStoreTrait};

// Index every trained entity embedding in an oxirs-vec store
let mut store = VectorStore::new();
for entity in model.get_entities() {
    let embedding = model.get_entity_embedding(&entity)?;
    store.insert_vector(entity, VecVector::new(embedding.values))?;
}

// Find the entities closest to Alice
let query = model.get_entity_embedding("http://example.org/Person/Alice")?;
let results = store.search_similar(&VecVector::new(query.values), 10)?;

With oxirs-chat (RAG)

oxirs-chat's rag module consumes entity/text embeddings produced by oxirs-embed as part of its retrieval-augmented generation pipeline. See oxirs-chat for the current pipeline API.

Performance

Benchmarks

Run the criterion suite for up-to-date numbers on your own hardware:

cargo bench -p oxirs-embed --bench performance_validation

Optimization Tips

use oxirs_embed::EmbeddingModel;

// encode() already batches: pass every text in one call instead of looping
let embeddings = model.encode(&texts).await?;

// Cache embeddings (fixed-capacity LRU, keyed by content hash + model id)
use oxirs_embed::embedding_cache::{CacheKey, EmbeddingCache};

let mut cache = EmbeddingCache::new(10_000);
let key = CacheKey::new(content_hash, "scibert");
if let Some(cached) = cache.get(&key) {
    // reuse cached embedding
} else {
    let embedding = model.encode(&[text.to_string()]).await?;
    cache.insert(key, embedding[0].clone());
}

GPU acceleration (feature gpu, Pure-Rust via scirs2-core/scirs2-linalg) is provided by oxirs_embed::gpu_acceleration::{GpuAccelerationConfig, GpuAccelerationManager} — see the crate-level docs for a full example.

Status

Production Release (v0.3.2)

• ✅ Structural KG embedding models: TransE, DistMult, ComplEx, RotatE, HoLE, ConvE, TuckER, QuatE
• ✅ Specialized domain text embeddings: SciBERT, CodeBERT, BioBERT, LegalBERT, FinBERT, ClinicalBERT, ChemBERT
• ✅ Link prediction, entity clustering, community detection, vector search, visualization, interpretability
• ✅ Entity linking and relation prediction (entity_linking module)
• ✅ Fine-tuning (fine_tuning/fine_tuner) and model ensembling (ensemble)
• ✅ Model zoo with SHA256-verified manifests (model_zoo)
• ✅ GPU acceleration behind the optional gpu feature (Pure-Rust, no CUDA/FFI required)
• 🚧 Contextual embeddings (contextual module) – implemented but not yet re-exported at the crate root while the API stabilizes

Contributing

This is an experimental module. Feedback welcome!

License

Apache-2.0

See Also

• oxirs-vec - Vector search engine
• oxirs-chat - AI-powered chat with RAG
• oxirs-core - RDF data model