feat: Make python tokenizer config pad strategy better (#344)

Before, `pad_strategy` was passed to `TokenizerConfig` as a str:

```python3
tokenizer_config = TokenizerConfig(pad_strategy="batch_longest")
```

This was nonfunctional, as `fixed` needs a number attached to it.
`fixed` means that the tokenizer always pads or truncates to a fixed
sequence length, and is used a lot in production deployments.

This pull request modifies the interface so that pad strategy options
are passed as classes instead of just a string value/enum.

For `batch_longest`:

```python3
from encoderfile import TokenizerConfig, BatchLongest, Fixed

# batch longest
tokenizer_config = TokenizerConfig(pad_strategy=BatchLongest())

# fixed
tokenizer_config = TokenizerConfig(pad_strategy=Fixed(512))
```

Author: besaleli

encoderfile-py/python/encoderfile/_core.pyi

@@ -37,9 +37,19 @@ class EncoderfileBuilder:
         no_download: bool = False,
     ): ...
 
+@final
+class BatchLongest:
+    pass
+
+@final
+class Fixed:
+    n: int
+
+    def __new__(cls, *, n: int) -> "Fixed": ...
+
 @final
 class TokenizerBuildConfig:
-    pad_strategy: Optional[str]
+    pad_strategy: Optional[BatchLongest | Fixed]
     truncation_side: Optional[str]
     truncation_strategy: Optional[str]
     max_length: Optional[int]
@@ -48,7 +58,7 @@ class TokenizerBuildConfig:
     def __new__(
         cls,
         *,
-        pad_strategy: Optional[str] = None,
+        pad_strategy: Optional[BatchLongest | Fixed] = None,
         truncation_side: Optional[str] = None,
         truncation_strategy: Optional[str] = None,
         max_length: Optional[int] = None,

encoderfile-py/python/encoderfile/enums.py

@@ -17,8 +17,3 @@ class TokenizerTruncationStrategy(StrEnum):
     LongestFirst = "longest_first"
     OnlyFirst = "only_first"
     OnlySecond = "only_second"
-
-
-class TokenizerPadStrategy(StrEnum):
-    BatchLongest = "batch_longest"
-    Fixed = "fixed"

encoderfile-py/src/builder.rs

@@ -13,12 +13,59 @@ use encoderfile::builder::{
     },
 };
 use encoderfile::common::{Config, ModelConfig};
+use pyo3::IntoPyObjectExt;
+use pyo3::exceptions::PyTypeError;
 use pyo3::types::PyString;
 use pyo3::{
     exceptions::{PyIOError, PyRuntimeError, PyValueError},
     prelude::*,
 };
 
+#[pyclass(name = "BatchLongest")]
+#[derive(Debug, Clone)]
+pub struct PyBatchLongest;
+
+#[pymethods]
+impl PyBatchLongest {
+    #[new]
+    pub fn new() -> Self {
+        PyBatchLongest
+    }
+}
+
+impl Default for PyBatchLongest {
+    fn default() -> Self {
+        PyBatchLongest
+    }
+}
+
+impl From<PyBatchLongest> for TokenizerPadStrategy {
+    fn from(_value: PyBatchLongest) -> Self {
+        TokenizerPadStrategy::BatchLongest
+    }
+}
+
+#[pyclass(name = "Fixed")]
+#[derive(Debug, Clone)]
+pub struct PyFixed {
+    #[pyo3(get, set)]
+    n: usize,
+}
+
+#[pymethods]
+impl PyFixed {
+    #[new]
+    pub fn new(n: usize) -> Self {
+        Self { n }
+    }
+}
+
+impl From<PyFixed> for TokenizerPadStrategy {
+    fn from(value: PyFixed) -> Self {
+        TokenizerPadStrategy::Fixed { fixed: value.n }
+    }
+}
+
 #[pyclass(name = "TokenizerBuildConfig", frozen)]
 pub struct PyTokenizerBuildConfig(TokenizerBuildConfig);
 
@@ -28,18 +75,27 @@ impl PyTokenizerBuildConfig {
     #[new]
     #[pyo3(signature = (*, pad_strategy = None, truncation_side = None, truncation_strategy = None, max_length = None, stride = None))]
     pub fn new(
-        pad_strategy: Option<String>,
+        pad_strategy: Option<Bound<'_, PyAny>>,
         truncation_side: Option<String>,
         truncation_strategy: Option<String>,
         max_length: Option<usize>,
         stride: Option<usize>,
     ) -> PyResult<Self> {
-        let pad_strategy = pad_strategy
-            .map(|s| {
-                s.parse::<TokenizerPadStrategy>()
-                    .map_err(|e| PyRuntimeError::new_err(format!("Failed to parse: {:?}", e)))
-            })
-            .transpose()?;
+        let pad_strategy = match pad_strategy {
+            Some(ps) => {
+                if let Ok(batch_longest) = ps.cast::<PyBatchLongest>() {
+                    Some(batch_longest.extract::<PyBatchLongest>().unwrap().into())
+                } else if let Ok(fixed) = ps.cast::<PyFixed>() {
+                    Some(fixed.extract::<PyFixed>().unwrap().into())
+                } else {
+                    return Err(PyTypeError::new_err(
+                        "Class must be BatchLongest, Fixed, or None",
+                    ));
+                }
+            }
+            None => None,
+        };
+
         let truncation_side = truncation_side
             .map(|s| {
                 s.parse::<TokenizerTruncationSide>()
@@ -63,8 +119,15 @@ impl PyTokenizerBuildConfig {
     }
 
     #[getter]
-    pub fn get_pad_strategy(&self) -> Option<String> {
-        self.0.pad_strategy.as_ref().map(|s| s.into())
+    pub fn get_pad_strategy(&self, py: Python<'_>) -> PyResult<Option<Py<PyAny>>> {
+        match self.0.pad_strategy.as_ref() {
+            Some(ps) => match ps {
+                TokenizerPadStrategy::BatchLongest => Some(PyBatchLongest::new().into_py_any(py)),
+                TokenizerPadStrategy::Fixed { fixed } => Some(PyFixed::new(*fixed).into_py_any(py)),
+            },
+            None => None,
+        }
+        .transpose()
     }
 
     #[getter]

encoderfile-py/src/lib.rs

@@ -23,6 +23,12 @@ mod encoderfile {
     #[pymodule_export]
     use super::builder::PyInspectInfo;
 
+    #[pymodule_export]
+    use super::builder::PyBatchLongest;
+
+    #[pymodule_export]
+    use super::builder::PyFixed;
+
     #[pymodule_export]
     use super::builder::PyTokenizerBuildConfig;
 }