car_inference/tasks/
classify.rs

1//! Classification — score text against candidate labels using prompt-based inference.
2
3use serde::{Deserialize, Serialize};
4
5#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
6use crate::backend::CandleBackend;
7#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
8use crate::tasks::generate;
9#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
10use crate::InferenceError;
11
12/// A classification request.
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct ClassifyRequest {
15    /// The text to classify.
16    pub text: String,
17    /// Candidate labels to score against.
18    pub labels: Vec<String>,
19    /// Optional model override.
20    pub model: Option<String>,
21}
22
23/// A classification result with label and confidence score.
24#[derive(Debug, Clone, Serialize, Deserialize)]
25pub struct ClassifyResult {
26    pub label: String,
27    pub score: f64,
28}
29
30#[cfg(not(all(target_os = "macos", target_arch = "aarch64", not(car_skip_mlx))))]
31/// Classify text against candidate labels.
32///
33/// Uses a prompt-based approach: asks the model to pick the best label,
34/// then parses the response. Falls back to first-token probability
35/// comparison when the response is ambiguous.
36pub async fn classify(
37    backend: &mut CandleBackend,
38    req: ClassifyRequest,
39) -> Result<Vec<ClassifyResult>, InferenceError> {
40    let labels_str = req
41        .labels
42        .iter()
43        .enumerate()
44        .map(|(i, l)| format!("{}. {}", i + 1, l))
45        .collect::<Vec<_>>()
46        .join("\n");
47
48    let prompt = format!(
49        "Classify the following text into one of these categories:\n\
50         {labels_str}\n\n\
51         Text: {}\n\n\
52         Respond with ONLY the category name, nothing else.",
53        req.text
54    );
55
56    let gen_req = generate::GenerateRequest {
57        prompt,
58        model: req.model.clone(),
59        params: generate::GenerateParams {
60            temperature: 0.0, // greedy for classification
61            max_tokens: 32,
62            ..Default::default()
63        },
64        context: None,
65        context_stable_prefix: None,
66        tools: None,
67        images: None,
68        messages: None,
69        cache_control: false,
70        response_format: None,
71        intent: None,
72        client_ref: None,
73        expected_row_digest: None,
74        expected_catalog_revision: None,
75        caller: None,
76    };
77
78    let (response, _ttft_ms, _prompt_tokens, _completion_tokens) =
79        generate::generate(backend, gen_req).await?;
80    let response_lower = response.trim().to_lowercase();
81
82    // Score each label based on string match in the response
83    let mut results: Vec<ClassifyResult> = req
84        .labels
85        .iter()
86        .map(|label| {
87            let label_lower = label.to_lowercase();
88            let score = if response_lower == label_lower {
89                1.0
90            } else if response_lower.contains(&label_lower) {
91                0.8
92            } else {
93                // Partial word overlap
94                let label_words: Vec<&str> = label_lower.split_whitespace().collect();
95                let matches = label_words
96                    .iter()
97                    .filter(|w| response_lower.contains(**w))
98                    .count();
99                if label_words.is_empty() {
100                    0.0
101                } else {
102                    0.5 * (matches as f64 / label_words.len() as f64)
103                }
104            };
105            ClassifyResult {
106                label: label.clone(),
107                score,
108            }
109        })
110        .collect();
111
112    // Sort by score descending
113    results.sort_by(|a, b| {
114        b.score
115            .partial_cmp(&a.score)
116            .unwrap_or(std::cmp::Ordering::Equal)
117    });
118
119    // Normalize scores to sum to 1
120    let total: f64 = results.iter().map(|r| r.score).sum();
121    if total > 0.0 {
122        for r in &mut results {
123            r.score /= total;
124        }
125    }
126
127    Ok(results)
128}
car_inference/tasks/classify.rs

car_inference/tasks/
classify.rs
