Add back OverwriteReason

rust/cli/Cargo.toml

@@ -2,7 +2,7 @@
 name = "magika-cli"
 version = "1.2.0-dev"
 license = "Apache-2.0"
-edition = "2021"
+edition = "2024"
 description = "Determines file content types using AI"
 repository = "https://github.com/google/magika"
 homepage = "https://securityresearch.google/magika"

rust/cli/src/backends.rs

@@ -23,7 +23,7 @@ use std::sync::{Arc, OnceLock};
 use anyhow::{Context, Result};
 use magika::{Backend, Builder, Runtime};
 
-use crate::{default_inference_threads, BackendChoice};
+use crate::{BackendChoice, default_inference_threads};
 
 /// Batches smaller than this stay on the CPU in automatic mode: a GPU only pays off on full ones.
 const GPU_MIN_BATCH: usize = 8;

rust/cli/src/main.rs

@@ -19,12 +19,12 @@ use std::io::{ErrorKind, Read, Write as _};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
-use anyhow::{ensure, Result};
+use anyhow::{Result, ensure};
 use clap::{Args, Parser, ValueEnum};
 use colored::ColoredString;
 use magika::{
-    self, Backend, ContentType, Features, FeaturesOrRuled, FileType, PredictionMode, Rules,
-    Runtime, TypeInfo,
+    self, Backend, ContentType, Features, FeaturesOrRuled, FileType, ModelDebugInfo,
+    OverwriteReason, Rules, Runtime, TypeInfo,
 };
 use serde::Serialize;
 
@@ -262,17 +262,6 @@ impl Stage {
     }
 }
 
-impl Flags {
-    fn debug_mode(&self) -> bool {
-        // We could try to optimize this and return false when possible.
-        true
-    }
-
-    fn prediction_mode(&self) -> PredictionMode {
-        PredictionMode::HighConfidence
-    }
-}
-
 /// Inference threads it takes to keep a GPU queued.
 ///
 /// The device is the bottleneck there, and it is already busy well before the host runs out of
@@ -325,8 +314,7 @@ fn main() -> Result<()> {
     let rules_only = flags.experimental.rules == RulesMode::Only;
     let mut builder = Runtime::builder();
     builder = builder.with_max_batch(batch_size);
-    builder.options.debug_mode = flags.debug_mode();
-    builder.options.prediction_mode = flags.prediction_mode();
+    builder.options.debug_mode = true;
     match flags.experimental.rules {
         RulesMode::Off => builder.options.use_builtin_rules = false,
         RulesMode::Enforce => (),
@@ -855,14 +843,24 @@ impl Response {
                     Some('s') => write!(&mut result, "{:.2}", self.score())?,
                     Some('S') => write!(&mut result, "{}%", (100. * self.score()).trunc())?,
                     Some('b') => {
-                        if let Some((inferred_type, score)) = self.best_guess(flags) {
-                            write!(
-                                &mut result,
-                                " [Low-confidence model best-guess: {} ({}), score={:.3}]",
-                                inferred_type.info().description,
-                                inferred_type.info().group,
-                                score,
-                            )?;
+                        if let Some(debug) = self.model_debug_info()
+                            && let Some(reason) = debug.overwrite_reason
+                        {
+                            let reason = match reason {
+                                OverwriteReason::LowConfidence => Some("Low-confidence"),
+                                OverwriteReason::OverwriteMap => None,
+                                OverwriteReason::RulesVeto => Some("Rules-veto"),
+                            };
+                            if let Some(reason) = reason {
+                                let best = debug.best_guess();
+                                let score = debug.output[best];
+                                let info = magika::MODEL_TRAINING_LABELS[best].info();
+                                write!(
+                                    &mut result,
+                                    " [{reason} model best-guess: {} ({}), score={score:.3}]",
+                                    info.description, info.group,
+                                )?;
+                            }
                         }
                     }
                     Some(c) => result.push(c),
@@ -875,24 +873,21 @@ impl Response {
         Ok(self.color(result.into()))
     }
 
-    // Returns the best guess if low confidence.
-    fn best_guess(&self, flags: &Flags) -> Option<(ContentType, f32)> {
-        let result = self.result.as_ref().ok()?;
-        let (label, score) = result.best_model_guess()?;
-        if flags.prediction_mode().is_confident(score, label) {
-            return None;
-        }
-        Some((magika::MODEL_TRAINING_LABELS[label], score))
+    fn model_debug_info(&self) -> Option<&ModelDebugInfo> {
+        self.result.as_ref().ok()?.debug_info()?.model()
     }
 
     fn json(self) -> Result<serde_json::Value> {
         let path = self.path.to_string_lossy();
+        let (score, dl) = match self.model_debug_info() {
+            Some(debug) => {
+                let best = debug.best_guess();
+                (debug.output[best], magika::MODEL_TRAINING_LABELS[best].info())
+            }
+            None => (1.0, ContentType::Undefined.info()),
+        };
         let result = match self.result {
             Ok(x) => {
-                let (score, dl) = match x.best_model_guess() {
-                    Some((label, score)) => (score, magika::MODEL_TRAINING_LABELS[label].info()),
-                    None => (1.0, ContentType::Undefined.info()),
-                };
                 let output = x.info();
                 let score = (score * 1000.).trunc() / 1000.;
                 let value = serde_json::to_value(JsonResult { dl, output, score })?;
@@ -939,10 +934,7 @@ impl Response {
     }
 
     fn score(&self) -> f32 {
-        match &self.result {
-            Err(_) => 1.0,
-            Ok(x) => x.best_model_guess().map_or(1.0, |x| x.1),
-        }
+        self.model_debug_info().map_or(1.0, |x| x.output[x.best_guess()])
     }
 
     fn color(&self, result: ColoredString) -> ColoredString {

rust/lib/src/file.rs

@@ -17,7 +17,7 @@ use std::borrow::Borrow;
 use anyhow::{Result, ensure};
 use ndarray::ArrayViewD;
 
-use crate::{ContentType, Features, MODEL_LABELS, Options, Rules};
+use crate::{ContentType, Features, MODEL_LABELS, MODEL_TRAINING_LABELS, Options, Rules};
 
 /// File types.
 ///
@@ -83,6 +83,22 @@ pub struct ModelDebugInfo {
     ///
     /// This list has the same length as [`MODEL_LABELS`] and is in one-to-one correspondance.
     pub output: Vec<f32>,
+
+    /// Reason why the inferred content type was overwritten.
+    pub overwrite_reason: Option<OverwriteReason>,
+}
+
+/// Reason to overwrite an inferred content type.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum OverwriteReason {
+    /// The inference score is too low for the inferred content type.
+    LowConfidence,
+
+    /// The inferred content type is not canonical.
+    OverwriteMap,
+
+    /// The rules vetoed the inferred content type.
+    RulesVeto,
 }
 
 impl FileType {
@@ -118,15 +134,29 @@ impl FileType {
         }
     }
 
-    /// Returns the best model guess.
-    ///
-    /// If debugging was enabled and the model triggered, returns the label index and score of the
-    /// best model guess.
-    pub fn best_model_guess(&self) -> Option<(usize, f32)> {
-        let FileType::Regular { debug_info: Some(x), .. } = self else { return None };
-        let DebugInfo::Model(x) = &**x else { return None };
-        let best = best_index(&x.output);
-        Some((best, x.output[best]))
+    /// Returns the debug information for regular files.
+    pub fn debug_info(&self) -> Option<&DebugInfo> {
+        match self {
+            FileType::Regular { debug_info, .. } => debug_info.as_deref(),
+            _ => None,
+        }
+    }
+}
+
+impl DebugInfo {
+    /// Returns the model debug information.
+    pub fn model(&self) -> Option<&ModelDebugInfo> {
+        match self {
+            DebugInfo::Model(x) => Some(x),
+            _ => None,
+        }
+    }
+}
+
+impl ModelDebugInfo {
+    /// Returns the index of the best model guess.
+    pub fn best_guess(&self) -> usize {
+        best_index(&self.output)
     }
 }
 
@@ -196,24 +226,30 @@ impl FileType {
             let builtin_rules_veto =
                 options.use_builtin_rules && Rules::builtin().vetos().contains(&overwrite);
             let rules_veto = custom_rules_veto || builtin_rules_veto;
-            let debug_info = options
-                .debug_mode
-                .then(|| Box::new(DebugInfo::Model(ModelDebugInfo { output: scores.to_vec() })));
-            let content_type = if ml_is_confident && !rules_veto {
-                overwrite
-            } else if overwrite.info().is_text {
-                ContentType::Txt
+            let fallback =
+                if overwrite.info().is_text { ContentType::Txt } else { ContentType::Unknown };
+            let (content_type, mut overwrite_reason) = if !ml_is_confident {
+                (fallback, Some(OverwriteReason::LowConfidence))
+            } else if rules_veto {
+                (fallback, Some(OverwriteReason::RulesVeto))
             } else {
-                ContentType::Unknown
+                (overwrite, Some(OverwriteReason::OverwriteMap))
             };
+            if content_type == MODEL_TRAINING_LABELS[best] {
+                overwrite_reason = None;
+            }
+            let debug_info = options.debug_mode.then(|| {
+                let output = scores.to_vec();
+                Box::new(DebugInfo::Model(ModelDebugInfo { output, overwrite_reason }))
+            });
             let content_types = vec![content_type];
             results.push(FileType::Regular { content_types, debug_info });
         }
         Ok(results)
     }
 }
 
-pub(crate) fn best_index(scores: &[f32]) -> usize {
+fn best_index(scores: &[f32]) -> usize {
     let mut best = 0;
     for (i, &x) in scores.iter().enumerate() {
         if scores[best].max(x) == x {

rust/lib/src/lib.rs

@@ -35,7 +35,9 @@
 pub use crate::backend::{Backend, BackendInfo};
 pub use crate::builder::Builder;
 pub use crate::content::{ContentType, MODEL_MAJOR_VERSION, MODEL_NAME};
-pub use crate::file::{DebugInfo, FileType, ModelDebugInfo, RulesDebugInfo, TypeInfo};
+pub use crate::file::{
+    DebugInfo, FileType, ModelDebugInfo, OverwriteReason, RulesDebugInfo, TypeInfo,
+};
 pub use crate::input::{Features, FeaturesOrRuled, Input};
 pub use crate::model::{MODEL_LABELS, MODEL_TRAINING_LABELS};
 pub use crate::options::{Options, PredictionMode};
@@ -109,32 +111,20 @@ mod tests {
         );
     }
 
-    fn model_info(
-        prediction_mode: PredictionMode, model: &ModelDebugInfo,
-    ) -> (f32, ContentType, ReferenceOverwriteReason) {
-        let best = crate::file::best_index(&model.output);
+    fn model_info(model: &ModelDebugInfo) -> (f32, ContentType, ReferenceOverwriteReason) {
+        let best = model.best_guess();
         let score = model.output[best];
         let dl = MODEL_TRAINING_LABELS[best];
-        let overwrite = MODEL_LABELS[best];
-        let overwrite_reason = if !prediction_mode.is_confident(score, best) {
-            let low_confidence =
-                if overwrite.info().is_text { ContentType::Txt } else { ContentType::Unknown };
-            if low_confidence == dl {
-                ReferenceOverwriteReason::None
-            } else {
-                ReferenceOverwriteReason::LowConfidence
-            }
-        } else if overwrite != dl {
-            ReferenceOverwriteReason::OverwriteMap
-        } else {
-            ReferenceOverwriteReason::None
+        let overwrite_reason = match model.overwrite_reason {
+            None => ReferenceOverwriteReason::None,
+            Some(OverwriteReason::LowConfidence) => ReferenceOverwriteReason::LowConfidence,
+            Some(OverwriteReason::OverwriteMap) => ReferenceOverwriteReason::OverwriteMap,
+            Some(OverwriteReason::RulesVeto) => unreachable!(),
         };
         (score, dl, overwrite_reason)
     }
 
-    fn assert_prediction(
-        prediction_mode: PredictionMode, actual: FileType, expected: Prediction, debug: &str,
-    ) {
+    fn assert_prediction(actual: FileType, expected: Prediction, debug: &str) {
         let FileType::Regular { content_types, debug_info: Some(debug_info) } = actual else {
             unreachable!()
         };
@@ -150,7 +140,7 @@ mod tests {
             DebugInfo::Model(x) => x,
             _ => unreachable!(),
         };
-        let (score, dl, overwrite_reason) = model_info(prediction_mode, model);
+        let (score, dl, overwrite_reason) = model_info(model);
         assert_float(score, expected.score, debug);
         assert_eq!(overwrite_reason, expected.overwrite_reason, "{debug}");
         assert_eq!(dl.info().label, expected.dl, "{debug}");
@@ -192,12 +182,11 @@ mod tests {
         let mut session = new_session();
         let mut checked = 0;
         for test in tests {
-            let prediction_mode = test.prediction_mode.into();
-            session.options_mut().prediction_mode = prediction_mode;
+            session.options_mut().prediction_mode = test.prediction_mode.into();
             assert_eq!(test.status, "ok"); // only scenario tested so far
             let expected = test.prediction.unwrap();
             let actual = session.identify_file(format!("../../{}", test.path)).unwrap();
-            assert_prediction(prediction_mode, actual, expected, &test.path);
+            assert_prediction(actual, expected, &test.path);
             checked += 1;
         }
         assert!(checked > 0, "reference fixture filter must exercise predictions");
@@ -222,13 +211,12 @@ mod tests {
         let mut session = new_session();
         let mut checked = 0;
         for test in tests {
-            let prediction_mode = test.prediction_mode.into();
-            session.options_mut().prediction_mode = prediction_mode;
+            session.options_mut().prediction_mode = test.prediction_mode.into();
             assert_eq!(test.status, "ok"); // only scenario tested so far
             let expected = test.prediction.unwrap();
             let content = BASE64.decode(test.content_base64.as_bytes()).unwrap();
             let actual = session.identify_content(content.as_slice()).unwrap();
-            assert_prediction(prediction_mode, actual, expected, &test.content_base64);
+            assert_prediction(actual, expected, &test.content_base64);
             checked += 1;
         }
         assert!(checked > 0, "reference fixture filter must exercise predictions");

rust/lib/src/tests/gpu.rs

@@ -70,7 +70,7 @@ fn reference_decisions_match_on_every_gpu_batch() {
             let expected = fixture.prediction.unwrap();
             match FeaturesOrRuled::extract_content(bytes.as_slice(), &options).unwrap() {
                 FeaturesOrRuled::Ruled(kind) => {
-                    assert_prediction(PredictionMode::HighConfidence, kind, expected, &name);
+                    assert_prediction(kind, expected, &name);
                 }
                 FeaturesOrRuled::Features(features) => {
                     let reference = cpu.run(&features.features, 1).unwrap();
@@ -135,7 +135,7 @@ fn reference_decisions_match_on_every_gpu_batch() {
                 let DebugInfo::Model(model) = &*debug_info else {
                     panic!("expected inference: {context}")
                 };
-                let (_, dl, reason) = model_info(PredictionMode::HighConfidence, model);
+                let (_, dl, reason) = model_info(model);
                 assert_eq!(content_type.info().label, sample.expected.output, "{context}");
                 assert_eq!(dl.info().label, sample.expected.dl, "{context}");
                 assert_eq!(reason, sample.expected.overwrite_reason, "{context}");