Use constructor in TargetSpec, allow union in builder (#325)

A normal constructor can now be used with `TargetSpec`. Also, the
builder config dict `target` accepts either a str or a prebuilt `Target
Spec`. The stubtest uses the right `_core` module now.

Since the `TargetSpec` is now part of the `EncoderfileConfig`, it needs
to derive `JsonSchema`. Let me know if that would be an issue. From the
top of my head I don't know if we currently test schema generation and
check something on it.

The `version` param takes `None` as default to make it consistent with
the Rust-side option. However, it might be more appropriate to signal
`'0.1.0'` as default, as it is the real value that will be set up when
the default is used. A discussion for another day 😅

Author: javiermtorres

Makefile

@@ -16,7 +16,7 @@ build-py:
 .PHONY: stubtest
 stubtest:
 	@echo "Running stubtest..."
-	@uv run --dev stubtest --allowlist=allowlist.txt encoderfile.encoderfile
+	@uv run --dev stubtest --allowlist=allowlist.txt encoderfile._core
 
 .PHONY: format
 format:

allowlist.txt

@@ -1,3 +1,3 @@
 # ignore missing __all__ in encoderfile
-encoderfile.encoderfile.__all__
+encoderfile._core.__all__
 # encoderfile.encoderfile.EncoderfileBuilder.from_config

…le-py/python/encoderfile/encoderfile.pyi …oderfile-py/python/encoderfile/_core.pyiencoderfile-py/python/encoderfile/encoderfile.pyi renamed to encoderfile-py/python/encoderfile/_core.pyi encoderfile-py/python/encoderfile/encoderfile.pyi renamed to encoderfile-py/python/encoderfile/_core.pyi

@@ -6,9 +6,7 @@ class TargetSpec:
     arch: str
     os: str
     abi: str
-
-    @staticmethod
-    def parse(spec: str) -> "TargetSpec": ...
+    def __new__(cls, spec: str): ...
 
 @final
 class EncoderfileBuilder:
@@ -20,7 +18,7 @@ class EncoderfileBuilder:
     def from_dict(
         *,
         name: str,
-        version: Optional[str] = "0.1.0",
+        version: Optional[str] = None,
         model_type: ModelType,
         path: str,
         output_path: Optional[str] = None,
@@ -30,7 +28,7 @@ class EncoderfileBuilder:
         lua_libs: Optional[list[str]] = None,
         tokenizer: Optional[TokenizerBuildConfig] = None,
         validate_transform: bool = True,
-        target: Optional[str] = None,
+        target: Optional[str | TargetSpec] = None,
     ) -> "EncoderfileBuilder": ...
     def build(
         self,

encoderfile-py/python/tests/test_encoderfile.py

@@ -197,6 +197,37 @@ def test_encoderfilebuilder_build_all_from_dict(tmp_path):
     assert result.encoderfile_config.lua_libs == config.get("lua_libs")
 
 
+def test_encoderfilebuilder_with_target_spec_object(tmp_path):
+    name = "model-built-from-config"
+    target_spec = TargetSpec("x86_64-unknown-linux-musl")
+    config = {
+        "name": name,
+        "path": "models/token_classification",
+        "model_type": ModelType.TokenClassification,
+        "output_path": str(tmp_path / f"{name}.encoderfile"),
+        "tokenizer": TokenizerBuildConfig.new(
+            pad_strategy="batch_longest",
+        ),
+        "transform": """
+        --- No docs
+        ---
+        --- 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
+        """,
+        "target": target_spec,
+    }
+    EncoderfileBuilder.from_dict(**config)
+
+
 def test_encoderfilebuilder_wrong_arch(tmp_path):
     name = "model-built-from-config"
     config = {
@@ -224,11 +255,10 @@ def test_encoderfilebuilder_wrong_arch(tmp_path):
         """,
         "target": "nonsense!",
     }
-    builder = EncoderfileBuilder.from_dict(**config)
-    with pytest.raises(RuntimeError) as exc_info:
-        builder.build(workdir=None, version=None, no_download=True)
+    with pytest.raises(ValueError) as exc_info:
+        EncoderfileBuilder.from_dict(**config)
     assert (
-        f"Error building encoderfile: invalid or unsupported target triple `{config['target']}`"
+        f"Failed to parse target spec: invalid or unsupported target triple `{config['target']}`"
         in str(exc_info.value)
     )
 
@@ -383,23 +413,23 @@ def test_tokenizer_config_partial_optional_fields():
 
 def test_parse_target_spec_valid_1():
     spec_str = "x86_64-unknown-linux-musl"
-    target_spec = TargetSpec.parse(spec_str)
+    target_spec = TargetSpec(spec_str)
     assert target_spec.arch == "x86_64"
     assert target_spec.os == "linux"
     assert target_spec.abi == "musl"
 
 
 def test_parse_target_spec_valid_2():
     spec_str = "aarch64-apple-darwin"
-    target_spec = TargetSpec.parse(spec_str)
+    target_spec = TargetSpec(spec_str)
     assert target_spec.arch == "aarch64"
     assert target_spec.os == "darwin"
     assert target_spec.abi == "gnu"
 
 
 def test_parse_target_spec_valid_3():
     spec_str = "x86_64-pc-windows-msvc"
-    target_spec = TargetSpec.parse(spec_str)
+    target_spec = TargetSpec(spec_str)
     assert target_spec.arch == "x86_64"
     assert target_spec.os == "windows"
     assert target_spec.abi == "msvc"
@@ -408,7 +438,7 @@ def test_parse_target_spec_valid_3():
 def test_parse_target_spec_invalid_format():
     spec_str = "nonsense!"
     with pytest.raises(ValueError) as exc_info:
-        TargetSpec.parse(spec_str)
+        TargetSpec(spec_str)
     assert (
         "Failed to parse target spec: invalid or unsupported target triple `nonsense!`"
         in str(exc_info.value)
@@ -418,7 +448,7 @@ def test_parse_target_spec_invalid_format():
 def test_parse_target_spec_unsupported_arch():
     spec_str = "riscv64-unknown-linux-musl"
     with pytest.raises(ValueError) as exc_info:
-        TargetSpec.parse(spec_str)
+        TargetSpec(spec_str)
     assert "Failed to parse target spec: unsupported architecture `riscv64`" in str(
         exc_info.value
     )

encoderfile-py/src/builder.rs

@@ -13,6 +13,7 @@ use encoderfile::builder::{
     },
 };
 use encoderfile::common::{Config, ModelConfig};
+use pyo3::types::PyString;
 use pyo3::{
     exceptions::{PyIOError, PyRuntimeError, PyValueError},
     prelude::*,
@@ -129,7 +130,7 @@ impl PyEncoderfileBuilder {
 
     #[allow(clippy::too_many_arguments)]
     #[staticmethod]
-    #[pyo3(signature = (*, name, version = Some("0.1.0"), model_type, path, output_path = None, cache_dir = None, base_binary_path = None, transform = None, lua_libs = None, tokenizer = None, validate_transform = true, target = None))]
+    #[pyo3(signature = (*, name, version = None, model_type, path, output_path = None, cache_dir = None, base_binary_path = None, transform = None, lua_libs = None, tokenizer = None, validate_transform = true, target = None))]
     fn from_dict(
         name: String,
         version: Option<&str>,
@@ -142,7 +143,7 @@ impl PyEncoderfileBuilder {
         lua_libs: Option<Vec<String>>,
         tokenizer: Option<Bound<'_, PyTokenizerBuildConfig>>,
         validate_transform: bool,
-        target: Option<String>,
+        target: Option<&Bound<PyAny>>,
     ) -> PyResult<Self> {
         let encoderfile = EncoderfileConfig {
             name,
@@ -158,7 +159,20 @@ impl PyEncoderfileBuilder {
             lua_libs: lua_libs.map(|libs| libs.into_iter().collect()),
             tokenizer: tokenizer.map(|t| t.borrow().0.clone()),
             validate_transform,
-            target,
+            target: target
+                .map(|t| {
+                    if let Ok(spec) = t.cast::<PyString>() {
+                        PyTargetSpec::parse(spec.to_str()?)
+                    } else if let Ok(spec) = t.cast::<PyTargetSpec>() {
+                        Ok(PyTargetSpec(spec.get().0.clone()))
+                    } else {
+                        Err(PyRuntimeError::new_err(
+                            "Failed to parse target spec: expected either a string or a TargetSpec",
+                        ))
+                    }
+                })
+                .transpose()?
+                .map(|t| t.0),
         };
         Ok(PyEncoderfileBuilder(EncoderfileBuilder {
             config: BuildConfig { encoderfile },
@@ -264,11 +278,12 @@ impl PyEncoderfileConfig {
 }
 
 #[pyclass(name = "TargetSpec", frozen)]
+#[derive(Clone)]
 pub struct PyTargetSpec(TargetSpec);
 
 #[pymethods]
 impl PyTargetSpec {
-    #[staticmethod]
+    #[new]
     #[pyo3(signature = (spec))]
     fn parse(spec: &str) -> PyResult<Self> {
         spec.parse()

encoderfile/src/builder/base_binary/target_spec.rs

@@ -1,8 +1,9 @@
 use anyhow::Result;
+use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use std::str::FromStr;
 
-#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash, JsonSchema)]
 pub struct TargetSpec {
     pub arch: Architecture,
     pub os: OperatingSystem,
@@ -125,7 +126,7 @@ impl std::fmt::Display for TargetSpec {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, JsonSchema)]
 pub enum Architecture {
     X86_64,
     Aarch64,
@@ -140,7 +141,7 @@ impl std::fmt::Display for Architecture {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, JsonSchema)]
 pub enum OperatingSystem {
     Linux,
     MacOS,
@@ -157,7 +158,7 @@ impl std::fmt::Display for OperatingSystem {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, JsonSchema)]
 pub enum Abi {
     Gnu,
     Musl,

encoderfile/src/builder/config.rs

@@ -54,15 +54,12 @@ pub struct EncoderfileConfig {
     pub tokenizer: Option<TokenizerBuildConfig>,
     #[serde(default = "default_validate_transform")]
     pub validate_transform: bool,
-    pub target: Option<String>,
+    pub target: Option<TargetSpec>,
 }
 
 impl EncoderfileConfig {
     pub fn target(&self) -> Result<Option<TargetSpec>> {
-        self.target
-            .as_ref()
-            .map(|s| TargetSpec::from_str(s.as_str()))
-            .transpose()
+        Ok(self.target.clone())
     }
 
     pub fn embedded_config(&self) -> Result<EmbeddedConfig> {