Add distribution to all embedders

This commit is contained in:
Louis Dureuil 2024-03-27 11:50:22 +01:00
parent 9a95ed619d
commit afd1da5642
No known key found for this signature in database
7 changed files with 41 additions and 18 deletions

View file

@ -33,6 +33,7 @@ enum WeightSource {
pub struct EmbedderOptions {
pub model: String,
pub revision: Option<String>,
pub distribution: Option<DistributionShift>,
}
impl EmbedderOptions {
@ -40,6 +41,7 @@ impl EmbedderOptions {
Self {
model: "BAAI/bge-base-en-v1.5".to_string(),
revision: Some("617ca489d9e86b49b8167676d8220688b99db36e".into()),
distribution: None,
}
}
}
@ -193,13 +195,15 @@ impl Embedder {
}
pub fn distribution(&self) -> Option<DistributionShift> {
if self.options.model == "BAAI/bge-base-en-v1.5" {
Some(DistributionShift {
current_mean: ordered_float::OrderedFloat(0.85),
current_sigma: ordered_float::OrderedFloat(0.1),
})
} else {
None
}
self.options.distribution.or_else(|| {
if self.options.model == "BAAI/bge-base-en-v1.5" {
Some(DistributionShift {
current_mean: ordered_float::OrderedFloat(0.85),
current_sigma: ordered_float::OrderedFloat(0.1),
})
} else {
None
}
})
}
}