feat: validate transforms on compile (#151)

- Adds compile-time checks for transforms during encoderfile build stage
- DRY mean pooling code a little

Scope creep:
- remove old config docs generation
- new makefile commands

Author: besaleli

.github/workflows/ci.yml

@@ -119,7 +119,7 @@ jobs:
         run: cargo binstall cargo-llvm-cov --force --no-confirm
 
       - name: Generate coverage
-        run: cargo llvm-cov --workspace --all-features --lcov --output-path lcov.info
+        run: make coverage
 
       - name: Upload to codecov
         uses: codecov/codecov-action@v2
@@ -163,7 +163,10 @@ jobs:
         with:
           toolchain: nightly
 
+      - name: Install test dependencies
+        run: make setup
+      
       - run: rustup component add --toolchain nightly-x86_64-unknown-linux-gnu clippy
 
       - name: Run Clippy
-        run: cargo clippy --all-features --all-targets -- -D warnings
+        run: make lint

.vscode/settings.json

@@ -1,5 +1,5 @@
 {
-  "rust-analyzer.cargo.features": ["bench"],
+  "rust-analyzer.cargo.features": ["dev-utils"],
   "rust-analyzer.procMacro.enable": true,
   "rust-analyzer.cargo.buildScripts.enable": true,
   "rust-analyzer.linkedProjects": [

Cargo.lock

@@ -12,6 +12,22 @@ format:
 	@echo "Formatting rust..."
 	@cargo fmt
 
+.PHONY: lint
+lint:
+	cargo clippy \
+		--all-features \
+		--all-targets \
+		-- \
+		-D warnings
+
+.PHONY: coverage
+coverage:
+	cargo llvm-cov \
+		--workspace \
+		--all-features \
+		--lcov \
+		--output-path lcov.info
+
 ,PHONY: licenses
 licenses:
 	@echo "Generating licenses..."
@@ -59,5 +75,3 @@ generate-docs:
 # 	generate JSON schema for encoderfile config
 	@cargo run \
 		--bin generate-encoderfile-config-schema
-# 	generate CLI docs for encoderfile build
-	@cargo run --bin generate-encoderfile-cli-docs --features="_internal"

Makefile

@@ -71,7 +71,10 @@ encoderfile:
   transform:
     path: ./transforms/normalize.lua
   # OR inline transform:
-  # transform: "return normalize(output)"
+  # transform: "function Postprocess(logits) return logits:lp_normalize(2.0, 2.0) end"
+
+  # Whether to validate transform with a dry-run (optional, defaults to true)
+  validate_transform: true
 
   # Whether to build the binary (optional, defaults to true)
   build: true

docs/reference/cli.md

@@ -1,3 +1,4 @@
+use encoderfile_core::transforms::Postprocessor;
 use ndarray::{Array2, Array3};
 use rand::Rng;
 
@@ -21,9 +22,9 @@ fn get_random_3d(x: usize, y: usize, z: usize) -> Array3<f32> {
 
 #[divan::bench(args = [(16, 16, 16), (32, 128, 384), (32, 256, 768)])]
 fn bench_embedding_l2_normalization(bencher: divan::Bencher, (x, y, z): (usize, usize, usize)) {
-    let engine = encoderfile_core::transforms::Transform::new(include_str!(
+    let engine = encoderfile_core::transforms::EmbeddingTransform::new(Some(include_str!(
         "../../transforms/embedding/l2_normalize_embeddings.lua"
-    ))
+    )))
     .unwrap();
 
     let test_tensor = get_random_3d(x, y, z);
@@ -35,8 +36,8 @@ fn bench_embedding_l2_normalization(bencher: divan::Bencher, (x, y, z): (usize,
 
 #[divan::bench(args = [(16, 2), (32, 8), (128, 32)])]
 fn bench_seq_cls_softmax(bencher: divan::Bencher, (x, y): (usize, usize)) {
-    let engine = encoderfile_core::transforms::Transform::new(include_str!(
-        "../../transforms/sequence_classification/softmax_logits.lua"
+    let engine = encoderfile_core::transforms::SequenceClassificationTransform::new(Some(
+        include_str!("../../transforms/sequence_classification/softmax_logits.lua"),
     ))
     .unwrap();
 
@@ -49,8 +50,8 @@ fn bench_seq_cls_softmax(bencher: divan::Bencher, (x, y): (usize, usize)) {
 
 #[divan::bench(args = [(16, 16, 2), (32, 128, 8), (128, 256, 32)])]
 fn bench_tok_cls_softmax(bencher: divan::Bencher, (x, y, z): (usize, usize, usize)) {
-    let engine = encoderfile_core::transforms::Transform::new(include_str!(
-        "../../transforms/token_classification/softmax_logits.lua"
+    let engine = encoderfile_core::transforms::TokenClassificationTransform::new(Some(
+        include_str!("../../transforms/token_classification/softmax_logits.lua"),
     ))
     .unwrap();
 

encoderfile-core/benches/benchmark_transforms.rs

@@ -3,18 +3,50 @@ use crate::{
         EmbeddingRequest, ModelType, SentenceEmbeddingRequest, SequenceClassificationRequest,
         TokenClassificationRequest,
     },
-    runtime::AppState,
+    runtime::{AppState, get_model, get_model_config, get_model_type, get_tokenizer},
     server::{run_grpc, run_http, run_mcp},
     services::{embedding, sentence_embedding, sequence_classification, token_classification},
 };
 use anyhow::Result;
+use clap::Parser;
 use clap_derive::{Parser, Subcommand, ValueEnum};
 use opentelemetry::trace::TracerProvider as _;
 use opentelemetry_otlp::{Protocol, WithExportConfig};
 use opentelemetry_sdk::trace::SdkTracerProvider;
 use std::{fmt::Display, io::Write};
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
+pub async fn cli_entrypoint(
+    model_bytes: &[u8],
+    config_str: &str,
+    tokenizer_json: &str,
+    model_type: &str,
+    model_id: &str,
+    transform_str: Option<&str>,
+) -> Result<()> {
+    let cli = Cli::parse();
+
+    let session = get_model(model_bytes);
+    let config = get_model_config(config_str);
+    let tokenizer = get_tokenizer(tokenizer_json, &config);
+    let model_type = get_model_type(model_type);
+    let transform_str = transform_str.map(|t| t.to_string());
+    let model_id = model_id.to_string();
+
+    let state = AppState {
+        session,
+        config,
+        tokenizer,
+        model_type,
+        model_id,
+        transform_str,
+    };
+
+    cli.command.execute(state).await?;
+
+    Ok(())
+}
+
 macro_rules! generate_cli_route {
     ($req:ident, $fn:path, $format:ident, $out_dir:expr, $state:expr) => {{
         let result = $fn($req, &$state)?;

encoderfile-core/src/cli.rs

@@ -1,4 +1,5 @@
 mod embedding;
+mod model_config;
 mod model_metadata;
 mod model_type;
 mod sentence_embedding;
@@ -7,6 +8,7 @@ mod token;
 mod token_classification;
 
 pub use embedding::*;
+pub use model_config::*;
 pub use model_metadata::*;
 pub use model_type::*;
 pub use sentence_embedding::*;

encoderfile-core/src/common/mod.rs

@@ -0,0 +1,87 @@
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ModelConfig {
+    pub model_type: String,
+    pub pad_token_id: u32,
+    pub num_labels: Option<usize>,
+    pub id2label: Option<HashMap<u32, String>>,
+    pub label2id: Option<HashMap<String, u32>>,
+}
+
+impl ModelConfig {
+    pub fn id2label(&self, id: u32) -> Option<&str> {
+        self.id2label.as_ref()?.get(&id).map(|s| s.as_str())
+    }
+
+    pub fn label2id(&self, label: &str) -> Option<u32> {
+        self.label2id.as_ref()?.get(label).copied()
+    }
+
+    pub fn num_labels(&self) -> Option<usize> {
+        if self.num_labels.is_some() {
+            return self.num_labels;
+        }
+
+        if let Some(id2label) = &self.id2label {
+            return Some(id2label.len());
+        }
+
+        if let Some(label2id) = &self.label2id {
+            return Some(label2id.len());
+        }
+
+        None
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_num_labels() {
+        let test_labels: Vec<(String, u32)> = vec![("a", 1), ("b", 2), ("c", 3)]
+            .into_iter()
+            .map(|(i, j)| (i.to_string(), j))
+            .collect();
+
+        let label2id: HashMap<String, u32> = test_labels.clone().into_iter().collect();
+        let id2label: HashMap<u32, String> = test_labels
+            .clone()
+            .into_iter()
+            .map(|(i, j)| (j, i))
+            .collect();
+
+        let config = ModelConfig {
+            model_type: "MyModel".to_string(),
+            pad_token_id: 0,
+            num_labels: Some(3),
+            id2label: Some(id2label.clone()),
+            label2id: Some(label2id.clone()),
+        };
+
+        assert_eq!(config.num_labels(), Some(3));
+
+        let config = ModelConfig {
+            model_type: "MyModel".to_string(),
+            pad_token_id: 0,
+            num_labels: None,
+            id2label: Some(id2label.clone()),
+            label2id: Some(label2id.clone()),
+        };
+
+        assert_eq!(config.num_labels(), Some(3));
+
+        let config = ModelConfig {
+            model_type: "MyModel".to_string(),
+            pad_token_id: 0,
+            num_labels: None,
+            id2label: None,
+            label2id: Some(label2id.clone()),
+        };
+
+        assert_eq!(config.num_labels(), Some(3));
+    }
+}

encoderfile-core/src/common/model_config.rs

@@ -1,8 +1,9 @@
 #[derive(Debug, Clone, serde::Serialize, utoipa::ToSchema)]
 #[serde(rename_all = "snake_case")]
+#[repr(u8)]
 pub enum ModelType {
-    Embedding,
-    SequenceClassification,
-    TokenClassification,
-    SentenceEmbedding,
+    Embedding = 1,
+    SequenceClassification = 2,
+    TokenClassification = 3,
+    SentenceEmbedding = 4,
 }