feat: add support for GLM-OCR model engine

- replace Quality enum with FormulaModel and TableModel options
- implement FormulaEngine and TableEngine for model dispatch
- add support for GLM-OCR in formula and table recognition
- update CLI and model loading to handle new engine selections

Author: mmgeorge

crates/papers-extract/src/bin/bench_formulas.rs

@@ -102,14 +102,18 @@ fn main() {
     let cache_dir = cli
         .model_cache_dir
         .unwrap_or_else(models::default_cache_dir);
-    let model_paths =
-        models::ensure_models(papers_extract::Quality::Fast, &cache_dir).expect("Model download");
+    let model_paths = models::ensure_models(
+        papers_extract::FormulaModel::PpFormulanet,
+        papers_extract::TableModel::SlanetPlus,
+        &cache_dir,
+    )
+    .expect("Model download");
 
     eprintln!("Loading formula predictor...");
     let predictor = FormulaPredictor::new(
-        &model_paths.formula_encoder,
-        &model_paths.formula_decoder,
-        &model_paths.formula_tokenizer,
+        model_paths.formula_encoder.as_deref().expect("formula_encoder path"),
+        model_paths.formula_decoder.as_deref().expect("formula_decoder path"),
+        model_paths.formula_tokenizer.as_deref().expect("formula_tokenizer path"),
     )
     .expect("FormulaPredictor::new");
     eprintln!("Formula predictor ready\n");

crates/papers-extract/src/bin/dump.rs

@@ -93,8 +93,12 @@ fn main() {
 
     // Load layout model
     eprintln!("Loading layout model...");
-    let paths = models::ensure_models(papers_extract::Quality::Fast, &cache_dir)
-        .expect("layout model files");
+    let paths = models::ensure_models(
+        papers_extract::FormulaModel::PpFormulanet,
+        papers_extract::TableModel::SlanetPlus,
+        &cache_dir,
+    )
+    .expect("layout model files");
     let layout = models::build_layout_detector(&paths.layout).expect("layout detector");
 
     // Load PDF

crates/papers-extract/src/lib.rs

@@ -25,8 +25,10 @@ pub struct ExtractOptions {
     pub confidence_threshold: f32,
     /// Whether to extract figures as images (default true).
     pub extract_images: bool,
-    /// Quality mode — affects model selection for tables (default Fast).
-    pub quality: Quality,
+    /// Formula recognition model (default PpFormulanet).
+    pub formula: FormulaModel,
+    /// Table recognition model (default SlanetPlus).
+    pub table: TableModel,
     /// Path to the pdfium binary (auto-detected if None).
     pub pdfium_path: Option<PathBuf>,
     /// Directory for ONNX model cache (auto-detected if None).
@@ -45,7 +47,8 @@ impl Default for ExtractOptions {
             dpi: 144,
             confidence_threshold: 0.3,
             extract_images: true,
-            quality: Quality::default(),
+            formula: FormulaModel::default(),
+            table: TableModel::default(),
             pdfium_path: None,
             model_cache_dir: None,
             page: None,
@@ -55,15 +58,26 @@ impl Default for ExtractOptions {
     }
 }
 
-/// Quality mode controls model selection for tables.
+/// Formula recognition model selection.
 #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
-pub enum Quality {
-    /// Fast mode (default): SLANet-Plus (7 MB).
+pub enum FormulaModel {
+    /// pp-formulanet split encoder/decoder (default).
     #[default]
-    Fast,
-    /// Quality mode: PP-LCNet table classifier (6.5 MB) + SLANeXt-wired (351 MB).
-    /// Better accuracy for complex tables.
-    Quality,
+    PpFormulanet,
+    /// GLM-OCR vision-language model.
+    GlmOcr,
+}
+
+/// Table recognition model selection.
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
+pub enum TableModel {
+    /// SLANet-Plus (7 MB, default).
+    #[default]
+    SlanetPlus,
+    /// PP-LCNet classifier + SLANeXt-wired (~358 MB).
+    SlanextWired,
+    /// GLM-OCR vision-language model with table prompt.
+    GlmOcr,
 }
 
 /// Controls what debug output to produce.

crates/papers-extract/src/main.rs

@@ -2,7 +2,7 @@ use std::path::PathBuf;
 use std::process;
 
 use clap::{Parser, ValueEnum};
-use papers_extract::{DebugMode, ExtractOptions, Quality};
+use papers_extract::{DebugMode, ExtractOptions, FormulaModel, TableModel};
 
 #[derive(Parser)]
 #[command(
@@ -18,9 +18,13 @@ struct Cli {
     #[arg(short, long)]
     output: Option<PathBuf>,
 
-    /// Quality mode for table recognition
-    #[arg(long, short, default_value = "fast")]
-    quality: QualityArg,
+    /// Formula recognition model
+    #[arg(long, default_value = "pp-formulanet")]
+    formula: FormulaArg,
+
+    /// Table recognition model
+    #[arg(long, default_value = "slanet-plus")]
+    table: TableArg,
 
     /// DPI for page rendering
     #[arg(long, default_value = "144")]
@@ -56,11 +60,21 @@ struct Cli {
 }
 
 #[derive(ValueEnum, Clone, Debug)]
-enum QualityArg {
+enum FormulaArg {
+    /// pp-formulanet split encoder/decoder
+    PpFormulanet,
+    /// GLM-OCR vision-language model
+    GlmOcr,
+}
+
+#[derive(ValueEnum, Clone, Debug)]
+enum TableArg {
     /// SLANet-Plus (7 MB)
-    Fast,
+    SlanetPlus,
     /// PP-LCNet classifier + SLANeXt-wired (~358 MB)
-    Quality,
+    SlanextWired,
+    /// GLM-OCR vision-language model
+    GlmOcr,
 }
 
 #[derive(ValueEnum, Clone, Debug)]
@@ -85,9 +99,14 @@ fn main() {
         dpi: cli.dpi,
         confidence_threshold: cli.confidence,
         extract_images: !cli.no_images,
-        quality: match cli.quality {
-            QualityArg::Fast => Quality::Fast,
-            QualityArg::Quality => Quality::Quality,
+        formula: match cli.formula {
+            FormulaArg::PpFormulanet => FormulaModel::PpFormulanet,
+            FormulaArg::GlmOcr => FormulaModel::GlmOcr,
+        },
+        table: match cli.table {
+            TableArg::SlanetPlus => TableModel::SlanetPlus,
+            TableArg::SlanextWired => TableModel::SlanextWired,
+            TableArg::GlmOcr => TableModel::GlmOcr,
         },
         pdfium_path: cli.pdfium_path,
         model_cache_dir: cli.model_cache_dir,

crates/papers-extract/src/models.rs

@@ -6,7 +6,7 @@ use oar_ocr::predictors::TableStructureRecognitionPredictor;
 use crate::error::ExtractError;
 use crate::formula::FormulaPredictor;
 use crate::glm_ocr::{GlmOcrConfig, GlmOcrPredictor};
-use crate::Quality;
+use crate::{FormulaModel, TableModel};
 
 /// Model file metadata for download.
 struct ModelFile {
@@ -90,14 +90,14 @@ const SLANEXT_WIRED: ModelFile = ModelFile {
 /// Resolved paths to all required model files.
 pub struct ModelPaths {
     pub layout: PathBuf,
-    pub slanet_plus: PathBuf,
-    pub table_dict: PathBuf,
-    pub formula_encoder: PathBuf,
-    pub formula_decoder: PathBuf,
-    pub formula_tokenizer: PathBuf,
-    // Quality-mode extras
+    pub slanet_plus: Option<PathBuf>,
+    pub table_dict: Option<PathBuf>,
+    pub formula_encoder: Option<PathBuf>,
+    pub formula_decoder: Option<PathBuf>,
+    pub formula_tokenizer: Option<PathBuf>,
     pub table_classifier: Option<PathBuf>,
     pub slanext_wired: Option<PathBuf>,
+    pub glm_ocr: Option<GlmOcrModelPaths>,
 }
 
 /// Resolved paths for GLM-OCR models (separate from pipeline models).
@@ -120,31 +120,50 @@ pub fn default_cache_dir() -> PathBuf {
         .join("models")
 }
 
-/// Ensure all models for the given quality mode are downloaded and return their paths.
+/// Ensure all models for the selected formula/table engines are downloaded.
 pub fn ensure_models(
-    quality: Quality,
+    formula: FormulaModel,
+    table: TableModel,
     cache_dir: &Path,
 ) -> Result<ModelPaths, ExtractError> {
     std::fs::create_dir_all(cache_dir)?;
 
-    // Always required
+    // Layout detection is always required
     let layout = ensure_model(cache_dir, &LAYOUT_MODEL)?;
-    let slanet_plus = ensure_model(cache_dir, &SLANET_PLUS)?;
-    let table_dict = ensure_model(cache_dir, &TABLE_DICT)?;
-    let formula_tokenizer = ensure_model(cache_dir, &FORMULA_TOKENIZER)?;
-
-    // Formula encoder/decoder (no auto-download, just check existence)
-    let formula_encoder = ensure_local_model(cache_dir, &FORMULA_ENCODER)?;
-    let formula_decoder = ensure_local_model(cache_dir, &FORMULA_DECODER)?;
-
-    // Table models selected by --quality
-    let (table_classifier, slanext_wired) = match quality {
-        Quality::Fast => (None, None),
-        Quality::Quality => {
+
+    // Formula models (only for pp-formulanet)
+    let (formula_encoder, formula_decoder, formula_tokenizer) = match formula {
+        FormulaModel::PpFormulanet => {
+            let enc = ensure_local_model(cache_dir, &FORMULA_ENCODER)?;
+            let dec = ensure_local_model(cache_dir, &FORMULA_DECODER)?;
+            let tok = ensure_model(cache_dir, &FORMULA_TOKENIZER)?;
+            (Some(enc), Some(dec), Some(tok))
+        }
+        FormulaModel::GlmOcr => (None, None, None),
+    };
+
+    // Table models (only for slanet variants)
+    let (slanet_plus, table_dict, table_classifier, slanext_wired) = match table {
+        TableModel::SlanetPlus => {
+            let slanet = ensure_model(cache_dir, &SLANET_PLUS)?;
+            let dict = ensure_model(cache_dir, &TABLE_DICT)?;
+            (Some(slanet), Some(dict), None, None)
+        }
+        TableModel::SlanextWired => {
+            let slanet = ensure_model(cache_dir, &SLANET_PLUS)?;
+            let dict = ensure_model(cache_dir, &TABLE_DICT)?;
             let classifier = ensure_model(cache_dir, &TABLE_CLASSIFIER)?;
             let wired = ensure_model(cache_dir, &SLANEXT_WIRED)?;
-            (Some(classifier), Some(wired))
+            (Some(slanet), Some(dict), Some(classifier), Some(wired))
         }
+        TableModel::GlmOcr => (None, None, None, None),
+    };
+
+    // GLM-OCR models (needed if either formula or table uses glm-ocr)
+    let glm_ocr = if formula == FormulaModel::GlmOcr || table == TableModel::GlmOcr {
+        Some(ensure_glm_ocr_models(cache_dir)?)
+    } else {
+        None
     };
 
     Ok(ModelPaths {
@@ -156,6 +175,7 @@ pub fn ensure_models(
         formula_tokenizer,
         table_classifier,
         slanext_wired,
+        glm_ocr,
     })
 }
 
@@ -393,21 +413,27 @@ fn platform_execution_providers() -> Vec<OrtExecutionProvider> {
 pub fn build_formula_predictor(
     paths: &ModelPaths,
 ) -> Result<FormulaPredictor, ExtractError> {
-    FormulaPredictor::new(
-        &paths.formula_encoder,
-        &paths.formula_decoder,
-        &paths.formula_tokenizer,
-    )
+    let enc = paths.formula_encoder.as_ref()
+        .ok_or_else(|| ExtractError::Model("formula_encoder path missing".into()))?;
+    let dec = paths.formula_decoder.as_ref()
+        .ok_or_else(|| ExtractError::Model("formula_decoder path missing".into()))?;
+    let tok = paths.formula_tokenizer.as_ref()
+        .ok_or_else(|| ExtractError::Model("formula_tokenizer path missing".into()))?;
+    FormulaPredictor::new(enc, dec, tok)
 }
 
 /// Build a standalone table structure recognition predictor.
 pub fn build_table_predictor(
     paths: &ModelPaths,
 ) -> Result<TableStructureRecognitionPredictor, ExtractError> {
+    let dict = paths.table_dict.as_ref()
+        .ok_or_else(|| ExtractError::Model("table_dict path missing".into()))?;
+    let slanet = paths.slanet_plus.as_ref()
+        .ok_or_else(|| ExtractError::Model("slanet_plus path missing".into()))?;
     let config = ort_config();
     TableStructureRecognitionPredictor::builder()
-        .dict_path(&paths.table_dict)
+        .dict_path(dict)
         .with_ort_config(config)
-        .build(&paths.slanet_plus)
+        .build(slanet)
         .map_err(|e| ExtractError::Model(format!("Failed to build table predictor: {e}")))
 }