Add inspect function to the python package (#306)

The `inspect` function has been added to the `encoderfile` package. The config classes have been annotated for python use via pyo3.

Author: javiermtorres

encoderfile-py/src/builder.rs

@@ -1,6 +1,10 @@
+use std::collections::HashMap;
 use std::path::PathBuf;
 
-use encoderfile::builder::builder::EncoderfileBuilder;
+use encoderfile::builder::cli::inspect::inspect_encoderfile;
+use encoderfile::builder::{builder::EncoderfileBuilder, cli::inspect::InspectInfo};
+use encoderfile::common::Config;
+use encoderfile::common::ModelConfig;
 use pyo3::{
     exceptions::{PyIOError, PyRuntimeError},
     prelude::*,
@@ -45,3 +49,82 @@ impl PyEncoderfileBuilder {
         Ok(())
     }
 }
+
+#[pyclass(name = "InspectInfo", frozen)]
+pub struct PyInspectInfo(InspectInfo);
+
+#[pymethods]
+impl PyInspectInfo {
+    #[getter]
+    fn get_model_config(&self) -> PyResult<PyModelConfig> {
+        Ok(PyModelConfig(self.0.model_config.clone()))
+    }
+    #[getter]
+    fn get_encoderfile_config(&self) -> PyResult<PyEncoderfileConfig> {
+        Ok(PyEncoderfileConfig(self.0.encoderfile_config.clone()))
+    }
+}
+
+#[pyclass(name = "ModelConfig", frozen)]
+pub struct PyModelConfig(ModelConfig);
+
+#[pymethods]
+impl PyModelConfig {
+    #[getter]
+    fn get_model_type(&self) -> String {
+        self.0.model_type.clone()
+    }
+
+    #[getter]
+    fn get_num_labels(&self) -> Option<usize> {
+        self.0.num_labels()
+    }
+
+    #[getter]
+    fn get_id2label(&self) -> Option<HashMap<u32, String>> {
+        self.0.id2label.clone()
+    }
+
+    #[getter]
+    fn get_label2id(&self) -> Option<HashMap<String, u32>> {
+        self.0.label2id.clone()
+    }
+}
+
+#[pyclass(name = "EncoderfileConfig", frozen)]
+pub struct PyEncoderfileConfig(Config);
+
+#[pymethods]
+impl PyEncoderfileConfig {
+    #[getter]
+    fn get_name(&self) -> String {
+        self.0.name.clone()
+    }
+
+    #[getter]
+    fn get_version(&self) -> String {
+        self.0.version.clone()
+    }
+
+    #[getter]
+    fn get_model_type(&self) -> String {
+        format!("{:?}", self.0.model_type)
+    }
+
+    #[getter]
+    fn get_transform(&self) -> Option<String> {
+        self.0.transform.clone()
+    }
+
+    #[getter]
+    fn get_lua_libs(&self) -> Option<Vec<String>> {
+        Some(self.0.lua_libs?.into())
+    }
+}
+
+#[pyfunction]
+pub fn inspect(_py: Python<'_>, path: &str) -> PyResult<PyInspectInfo> {
+    let result = inspect_encoderfile(path)
+        .map_err(|e| PyRuntimeError::new_err(format!("Failed to inspect encoderfile: {:?}", e)))?;
+    Ok(PyInspectInfo(result))
+}

encoderfile-py/src/lib.rs

@@ -7,4 +7,7 @@ mod builder;
 mod encoderfile {
     #[pymodule_export]
     use super::builder::PyEncoderfileBuilder;
+
+    #[pymodule_export]
+    use super::builder::inspect;
 }

encoderfile/src/builder/cli/inspect.rs

@@ -14,22 +14,26 @@ use crate::{
 
 // inspect struct with info
 
-#[derive(Debug, Serialize)]
+#[derive(Debug, Serialize, Clone)]
 pub struct InspectInfo {
     pub model_config: ModelConfig,
     pub encoderfile_config: Config,
 }
 
-pub fn inspect_encoderfile(path_str: &String) -> Result<String> {
+pub fn inspect_encoderfile(path_str: &str) -> Result<InspectInfo> {
     let file = File::open(Path::new(&path_str))?;
     let mut file = BufReader::new(file);
     let mut loader = load_assets(&mut file)?;
 
     let config = loader.encoderfile_config()?;
     let model_config = loader.model_config()?;
 
-    Ok(to_string_pretty(&InspectInfo {
+    Ok(InspectInfo {
         model_config,
         encoderfile_config: config,
-    })?)
+    })
+}
+
+pub fn inspect_encoderfile_pretty(path_str: &str) -> Result<String> {
+    Ok(to_string_pretty(&inspect_encoderfile(path_str)?)?)
 }

encoderfile/src/builder/cli/mod.rs

@@ -4,7 +4,7 @@ use std::path::PathBuf;
 use clap_derive::{Args, Parser, Subcommand};
 
 mod build;
-mod inspect;
+pub mod inspect;
 mod runtime;
 
 #[cfg(feature = "dev-utils")]
@@ -70,7 +70,7 @@ impl Commands {
             Self::Runtime(r) => r.execute(global),
             Self::NewTransform { model_type } => super::transforms::new_transform(model_type),
             Self::Inspect { path } => {
-                println!("{}", inspect::inspect_encoderfile(&path)?);
+                println!("{}", inspect::inspect_encoderfile_pretty(&path)?);
                 Ok(())
             }
         }

encoderfile/src/common/config.rs

@@ -3,7 +3,7 @@ use anyhow::{Result, bail};
 use serde::{Deserialize, Serialize};
 use tokenizers::{PaddingParams, TruncationParams};
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct Config {
     pub name: String,
     pub version: String,
@@ -57,6 +57,26 @@ impl TryFrom<Vec<String>> for LuaLibs {
     }
 }
 
+impl From<LuaLibs> for Vec<String> {
+    fn from(value: LuaLibs) -> Vec<String> {
+        [
+            ("coroutine", value.coroutine),
+            ("table", value.table),
+            ("io", value.io),
+            ("os", value.os),
+            ("string", value.string),
+            ("utf8", value.utf8),
+            ("math", value.math),
+            ("package", value.package),
+            ("debug", value.debug),
+        ]
+        .into_iter()
+        .filter(|&(_k, v)| v)
+        .map(|(k, _)| k.to_string())
+        .collect()
+    }
+}
+
 #[derive(Debug, Clone, Default, Serialize, Deserialize)]
 pub struct TokenizerConfig {
     pub padding: PaddingParams,

encoderfile/src/common/model_config.rs

@@ -1,7 +1,7 @@
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct ModelConfig {
     pub model_type: String,
     pub num_labels: Option<usize>,

test_config.yml

@@ -2,7 +2,7 @@ encoderfile:
   name: my-model-2
   path: models/token_classification
   model_type: token_classification
-  output_path: ./my-model-2.encoderfile
+  output_path: ./test-model.encoderfile
   transform: |
     --- Applies a softmax across token classification logits.
     --- Each token classification is normalized independently.

test_config_lua.yml

@@ -0,0 +1,23 @@
+encoderfile:
+  name: my-model-2-lua
+  path: models/dummy_electra_token_classifier
+  model_type: token_classification
+  output_path: ./test-model-lua.encoderfile
+  lua_libs:
+    - table
+    - math
+  transform: |
+    --- Applies a softmax across token classification logits.
+    --- Each token classification is normalized independently.
+    --- 
+    --- Args:
+    ---   arr (Tensor): A tensor of shape [batch_size, n_tokens, n_labels].
+    ---                 The softmax is applied along the third axis (n_labels).
+    ---
+    --- Returns:
+    ---   Tensor: The input tensor with softmax-normalized embeddings.
+    ---@param arr Tensor
+    ---@return Tensor
+    function Postprocess(arr)
+        return arr:softmax(3)
+    end