# Embedding Trainer
[](https://opensource.org/licenses/Apache-2.0)
[](https://www.rust-lang.org/)
[](https://github.com/yourusername/embedding-trainer)
A fast and flexible Rust library and CLI tool for training word embeddings from scratch using Skip-gram and CBOW algorithms with built-in validation and evaluation.
## ✨ Features
### 🚀 **Algorithms**
- **Skip-gram**: Predicts context words given target words
- **CBOW**: Predicts target words given context words
### 📊 **Training Features**
- Configurable embedding dimensions
- Adjustable learning rates and epochs
- Customizable context windows
- Negative sampling support
- Batch processing capabilities
### 🔧 **CLI Tools**
- **Training**: Train embeddings from text data with optional validation split
- **Similarity**: Calculate semantic similarity between words
- **Inspection**: Analyze trained models and vocabulary
- **Export**: Save embeddings in multiple formats (text, JSON, binary, Word2Vec)
- **Validate**: Evaluate a saved model on held-out validation text
### 💾 **Data Support**
- Text file processing
- Vocabulary management
- Model persistence
- Multiple export formats
- Streaming support for large datasets
## 🚀 Quick Start
### Installation
```bash
# Clone the repository
git clone https://github.com/yourusername/embedding-trainer.git
cd embedding-trainer
# Build the project
cargo build --release
# Or install locally
cargo install --path .
```
### Basic Usage
#### 1. Train Your First Embeddings
```bash
# Prepare your training data
echo "the quick brown fox jumps over the lazy dog" > data.txt
# Train embeddings using Skip-gram
embedding train \
--input data.txt \
--output model.json \
--embeddings embeddings.txt \
--dim 100 \
--epochs 10 \
--model-type skipgram
# Train with validation split
embedding train \
--input data.txt \
--output model.json \
--embeddings embeddings.txt \
--dim 100 \
--epochs 10 \
--validation-ratio 0.2 \
--validation-output metrics.json
```
#### 2. Calculate Similarity
```bash
# Calculate similarity between words
embedding similarity fox dog \
--model model.json
# Expected output:
# Similarity between 'fox' and 'dog': 0.8234
```
#### 3. Inspect Model
```bash
# View model information
embedding info --model model.json
# Shows vocabulary size, embedding dimension, training config
```
#### 4. Export Embeddings
```bash
# Export to different formats
embedding export \
--model model.json \
--output embeddings.json \
--format json
```
## 📚 Library Usage
### Basic Example
```rust
use embedding::*;
fn main() -> Result<(), String> {
// Load and prepare data
let text = "the quick brown fox jumps over the lazy dog";
let sentences = load_text_data(text);
let (vocab, reverse_vocab) = build_vocab(&sentences);
let training_data = TrainingData {
sentences,
vocab,
reverse_vocab,
};
// Configure training
let config = TrainingConfig {
embedding_dim: 300,
learning_rate: 0.025,
epochs: 10,
batch_size: 32,
context_window: 5,
negative_samples: 5,
model_type: ModelType::SkipGram,
lr_schedule: LearningRateSchedule::Constant,
early_stopping: None,
l2_regularization: None,
gradient_clip: None,
validation_ratio: None,
};
// Train model
let mut model = EmbeddingModel::new(config, training_data.vocab.len());
model.train(&training_data)?;
// Calculate similarity
if let Some(similarity) = model.similarity("fox", "dog", &training_data) {
println!("Similarity: {:.4}", similarity);
}
// Save model
model.save_embeddings("embeddings.txt", &training_data)?;
Ok(())
}
```
### Advanced Usage
```rust
use embedding::*;
use std::fs;
fn advanced_example() -> Result<(), String> {
// Load large dataset with streaming
let text = fs::read_to_string("large_dataset.txt")?;
let sentences = load_text_data(&text);
// Build vocabulary with size limit
let (vocab, reverse_vocab) = build_vocab(&sentences);
println!("Vocabulary size: {}", vocab.len());
let training_data = TrainingData {
sentences,
vocab,
reverse_vocab,
};
// Configure advanced training parameters
let config = TrainingConfig {
embedding_dim: 500,
learning_rate: 0.01,
epochs: 50,
batch_size: 128,
context_window: 10,
negative_samples: 10,
model_type: ModelType::Cbow,
lr_schedule: LearningRateSchedule::Constant,
early_stopping: None,
l2_regularization: None,
gradient_clip: None,
validation_ratio: None,
};
// Train with multiple epochs
let mut model = EmbeddingModel::new(config, training_data.vocab.len());
// Train in chunks for large datasets
for epoch in 0..10 {
println!("Training epoch {}/10", epoch + 1);
model.train(&training_data)?;
}
// Export to multiple formats
model.save_embeddings("embeddings.txt", &training_data)?;
println!("Training completed!");
Ok(())
}
```
## 🔧 Configuration
### Training Parameters
| `--dim` | Embedding dimension | 300 | 10-1000 |
| `--learning-rate` | Learning rate | 0.025 | 0.001-1.0 |
| `--epochs` | Number of training epochs | 10 | 1-1000 |
| `--batch-size` | Mini-batch size | 32 | 1-1000 |
| `--window` | Context window size | 5 | 1-20 |
| `--negative-samples` | Number of negative samples | 5 | 1-20 |
| `--validation-ratio` | Fraction of data for validation | 0.0 | 0.0-0.5 |
| `--validation-output` | File to write validation metrics JSON | - | - |
### Algorithm Types
- **`skipgram`**: Skip-gram algorithm (default)
- **`cbow`**: Continuous Bag of Words
### Export Formats
- **`text`**: Plain text format (default)
- **`json`**: JSON format with metadata
- **`bin`**: Binary format using bincode
- **`word2vec`**: Word2Vec/Gensim text format
## 📖 CLI Reference
### Training Command
```bash
embedding train [OPTIONS]
```
**Options:**
- `--input <FILE>` - Input text file (required)
- `--output <FILE>` - Output model file (required)
- `--embeddings <FILE>` - Embeddings output file (required)
- `--dim <SIZE>` - Embedding dimension (default: 300)
- `--learning-rate <RATE>` - Learning rate (default: 0.025)
- `--epochs <COUNT>` - Number of epochs (default: 10)
- `--batch-size <SIZE>` - Batch size (default: 32)
- `--window <SIZE>` - Context window size (default: 5)
- `--negative-samples <COUNT>` - Negative samples (default: 5)
- `--model-type <TYPE>` - Algorithm type (skipgram|cbow)
- `--validation-ratio <RATIO>` - Fraction for validation (default: 0.0)
- `--validation-output <FILE>` - Path to write validation metrics JSON
### Similarity Command
```bash
embedding similarity <WORD1> <WORD2> [OPTIONS]
```
**Options:**
- `--model <FILE>` - Model file (required)
### Info Command
```bash
embedding info [OPTIONS]
```
**Options:**
- `--model <FILE>` - Model file (required)
### Export Command
```bash
embedding export [OPTIONS]
```
**Options:**
- `--model <FILE>` - Model file (required)
- `--output <FILE>` - Output file (required)
- `--format <FORMAT>` - Export format (text|json|bin|word2vec)
### Validate Command
```bash
embedding validate [OPTIONS]
```
**Options:**
- `--model <FILE>` - Model file (required)
- `--input <FILE>` - Validation text file (required)
- `--output <FILE>` - Output metrics JSON file (optional)
## 🔍 Examples
### Example 1: Basic Word Embeddings
```bash
# Create sample data
cat > animals.txt << EOF
cat meows loudly
dog barks loudly
bird sings beautifully
fish swims quietly
horse gallops fast
EOF
# Train embeddings
embedding train \
--input animals.txt \
--output animal_model.json \
--embeddings animal_embeddings.txt \
--dim 50 \
--epochs 20 \
--model-type skipgram
# Test similarity
embedding similarity cat dog \
--model animal_model.json
```
### Example 2: Document Embeddings
```bash
# Prepare document data
cat > documents.txt << EOF
Machine learning is a subset of artificial intelligence.
Deep learning uses neural networks with multiple layers.
Natural language processing deals with text and speech.
Computer vision enables computers to understand images.
EOF
# Train with CBOW and validation
embedding train \
--input documents.txt \
--output doc_model.json \
--embeddings doc_embeddings.txt \
--dim 100 \
--epochs 15 \
--model-type cbow \
--validation-ratio 0.2
```
### Example 3: Large Dataset Processing
```bash
# Process large file with multiple epochs
embedding train \
--input large_corpus.txt \
--output large_model.json \
--embeddings large_embeddings.txt \
--dim 300 \
--epochs 50 \
--batch-size 256 \
--window 10 \
--model-type cbow
```
## 🧪 Development
### Building from Source
```bash
# Clone repository
git clone https://github.com/yourusername/embedding-trainer.git
cd embedding-trainer
# Build development version
cargo build
# Run tests
cargo test
# Run benchmarks
cargo bench
# Build documentation
cargo doc --open
```
### Running Tests
```bash
# Run all tests
cargo test
# Run specific test
cargo test test_build_vocab
# Run with verbose output
cargo test -- --verbose
```
### Development Features
- **Unit Tests**: Comprehensive test coverage
- **Integration Tests**: End-to-end testing
- **Benchmarks**: Performance testing
- **Documentation**: API documentation
## 📊 Performance
### Benchmarks
| Skip-gram | 10K words | 300 | 2.3s | 45MB |
| CBOW | 10K words | 300 | 1.8s | 42MB |
### Optimization Tips
1. **Use appropriate batch sizes** for your dataset
2. **Adjust learning rate** based on dataset size
3. **Context window size** affects training speed and quality
4. **Use negative sampling** for large vocabularies
5. **Monitor memory usage** with large datasets
## 🤝 Contributing
We welcome contributions! Please see our [Contributing Guide](CONTRIBUTING.md) for details.
### Development Workflow
1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests for new functionality
5. Run the test suite
6. Submit a pull request
### Code Style
- Follow Rust formatting standards
- Use `cargo fmt` for code formatting
- Add comprehensive documentation
- Include tests for new features
## 📈 Roadmap
### Version 1.0 (Current)
- ✅ Skip-gram and CBOW algorithms
- ✅ CLI interface with train, validate, similarity, info, export
- ✅ Model persistence (JSON, binary, Word2Vec, ONNX, NumPy)
- ✅ Similarity calculations and semantic search
- ✅ Validation split and evaluation metrics (accuracy, precision, recall, F1)
- ✅ Learning rate scheduling (constant, exponential, step, cosine)
- ✅ Early stopping and L2 regularization
### Version 1.1 (Planned)
- GPU acceleration
- Advanced tokenization improvements
- Cross-validation support
- Learning curve visualization
### Version 2.0 (Future)
- Transformer-based models
- Multi-modal embeddings
- Real-time training
- Standard word similarity benchmarks integration
## 🐛 Troubleshooting
### Common Issues
1. **Memory Error with Large Datasets**
- Reduce batch size
- Use streaming processing
- Increase system memory
2. **Poor Similarity Results**
- Increase training epochs
- Adjust learning rate
- Try different algorithms
3. **Missing Words in Vocabulary**
- Check text preprocessing
- Verify tokenization
- Ensure words appear in text
### Performance Issues
- **Slow Training**: Reduce batch size or use negative sampling
- **High Memory Usage**: Use smaller embedding dimensions
- **Poor Quality**: Increase epochs or adjust parameters
## 📄 License
This project is licensed under the Apache License 2.0 - see the [LICENSE](LICENSE) file for details.
## 🙏 Acknowledgments
- Inspired by Word2Vec, GloVe, and BERT
- Built with [ndarray](https://github.com/rust-ndarray/ndarray) for numerical computing
- CLI powered by [clap](https://github.com/clap-rs/clap)
- Serialization using [serde](https://serde.rs/)
## 📞 Support
- 📧 **Email**: your.email@example.com
- 💬 **Discussions**: [GitHub Discussions](https://github.com/yourusername/embedding-trainer/discussions)
- 🐛 **Issues**: [GitHub Issues](https://github.com/yourusername/embedding-trainer/issues)
- 📖 **Documentation**: [docs.rs/embedding-trainer](https://docs.rs/embedding-trainer)
---
**Made with ❤️ by the Embedding Trainer Team**
*For the latest updates, check our [GitHub repository](https://github.com/yourusername/embedding-trainer)*