Struct BoxEmbeddingTrainer 

pub struct BoxEmbeddingTrainer { /* private fields */ }

Trainer for box embedding models.

Implementations

impl BoxEmbeddingTrainer

pub fn new( config: TrainingConfig, dim: usize, initial_embeddings: Option<HashMap<usize, Vec<f32>>>, ) -> Self

Create a new trainer.

Arguments

• config - Training configuration
• dim - Embedding dimension
• initial_embeddings - Optional initial vector embeddings (entity_id → vector)

pub fn initialize_boxes( &mut self, examples: &[TrainingExample], initial_embeddings: Option<&HashMap<usize, Vec<f32>>>, )

Initialize boxes from entities.

Creates trainable boxes for all entities, either from provided vector embeddings or random initialization.

Key insight: For positive pairs (entities that corefer), initialize boxes to overlap so gradients can flow from the start.

Arguments

• examples - Training examples with entities and coreference chains
• initial_embeddings - Optional pre-computed vector embeddings (entity_id → vector) If provided, boxes are initialized around these vectors. If None, uses smart random initialization with shared centers for coreferent entities.

pub fn train(&mut self, examples: &[TrainingExample]) -> Vec<f32>

Train on a dataset with mini-batching and early stopping. Uses adaptive negative weighting: starts with low weight to learn positives, then gradually increases to separate negatives.

pub fn get_boxes(&self) -> HashMap<usize, BoxEmbedding>

Get trained boxes for inference.

pub fn get_overlap_stats(&self, examples: &[TrainingExample]) -> (f32, f32, f32)

Get diagnostic statistics about box overlaps.

Returns (avg_positive_score, avg_negative_score, overlap_rate)

pub fn evaluate( &self, examples: &[TrainingExample], threshold: f32, ) -> (f32, f32, f32, f32)

Evaluate coreference accuracy on a test set.

Returns (accuracy, precision, recall, f1) where:

• Accuracy: fraction of pairs correctly classified
• Precision: fraction of predicted positives that are correct
• Recall: fraction of true positives that are predicted
• F1: harmonic mean of precision and recall

Note: This is a simple pair-wise evaluation. For standard coreference metrics (MUC, B³, CEAF, LEA, BLANC, CoNLL F1), use evaluate_standard_metrics() instead.

pub fn save_boxes(&self, path: &str) -> Result<(), Box<dyn Error>>

Save trained boxes to a file (JSON format).

Arguments

• path - File path to save to

Example

trainer.save_boxes("trained_boxes.json")?;

pub fn load_boxes( path: &str, dim: usize, ) -> Result<HashMap<usize, TrainableBox>, Box<dyn Error>>

Load trained boxes from a file (JSON format).

Arguments

• path - File path to load from
• dim - Expected embedding dimension

Example

let boxes = BoxEmbeddingTrainer::load_boxes("trained_boxes.json", 32)?;