use finetype_core::tokenizer::TokenizerError;
use thiserror::Error;
#[derive(Error, Debug)]
pub enum InferenceError {
#[error("Model error: {0}")]
ModelError(#[from] candle_core::Error),
#[error("Tokenizer error: {0}")]
TokenizerError(#[from] TokenizerError),
#[error("Taxonomy error: {0}")]
TaxonomyError(#[from] finetype_core::taxonomy::TaxonomyError),
#[error("IO error: {0}")]
IoError(#[from] std::io::Error),
#[error("Invalid model path: {0}")]
InvalidPath(String),
}
#[derive(Debug, Clone)]
pub struct ClassificationResult {
pub label: String,
pub confidence: f32,
pub all_scores: Vec<(String, f32)>,
}
pub trait ValueClassifier: Send + Sync {
fn classify(&self, text: &str) -> Result<ClassificationResult, InferenceError>;
fn classify_batch(&self, texts: &[String])
-> Result<Vec<ClassificationResult>, InferenceError>;
}
#[cfg(test)]
pub struct MockClassifier {
label: String,
}
#[cfg(test)]
impl MockClassifier {
pub fn new(label: &str) -> Self {
Self {
label: label.to_string(),
}
}
}
#[cfg(test)]
impl ValueClassifier for MockClassifier {
fn classify(&self, _text: &str) -> Result<ClassificationResult, InferenceError> {
Ok(ClassificationResult {
label: self.label.clone(),
confidence: 0.8,
all_scores: vec![(self.label.clone(), 0.8)],
})
}
fn classify_batch(
&self,
texts: &[String],
) -> Result<Vec<ClassificationResult>, InferenceError> {
Ok(texts
.iter()
.map(|_| ClassificationResult {
label: self.label.clone(),
confidence: 0.8,
all_scores: vec![(self.label.clone(), 0.8)],
})
.collect())
}
}