use super::types::{EmbeddingConfig, TransformerModelType};
use super::sentencetransformergenerator_type::SentenceTransformerGenerator;
impl SentenceTransformerGenerator {
pub fn with_model_type(config: EmbeddingConfig, model_type: TransformerModelType) -> Self {
Self { config, model_type }
}
pub fn roberta(config: EmbeddingConfig) -> Self {
Self::with_model_type(config, TransformerModelType::RoBERTa)
}
pub fn distilbert(config: EmbeddingConfig) -> Self {
let adjusted_config = EmbeddingConfig {
dimensions: 384,
..config
};
Self::with_model_type(adjusted_config, TransformerModelType::DistilBERT)
}
pub fn multilingual_bert(config: EmbeddingConfig) -> Self {
Self::with_model_type(config, TransformerModelType::MultiBERT)
}
}