Fixing missing direction in TruncationParams. (#868)

huggingface · Jan 4, 2022 · 4122a33 · 4122a33
1 parent 7069988
commit 4122a33
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 14 deletions.
diff --git a/bindings/python/test.py b/bindings/python/test.py
diff --git a/tokenizers/src/utils/truncation.rs b/tokenizers/src/utils/truncation.rs
@@ -8,6 +8,11 @@ pub enum TruncationDirection {
     Left,
     Right,
 }
+impl Default for TruncationDirection {
+    fn default() -> Self {
+        TruncationDirection::Right
+    }
+}
 
 impl std::convert::AsRef<str> for TruncationDirection {
     fn as_ref(&self) -> &str {
@@ -20,6 +25,7 @@ impl std::convert::AsRef<str> for TruncationDirection {
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct TruncationParams {
+    #[serde(default)]
     pub direction: TruncationDirection,
     pub max_length: usize,
     pub strategy: TruncationStrategy,
@@ -30,9 +36,9 @@ impl Default for TruncationParams {
     fn default() -> Self {
         Self {
             max_length: 512,
-            strategy: TruncationStrategy::LongestFirst,
+            strategy: TruncationStrategy::default(),
             stride: 0,
-            direction: TruncationDirection::Right,
+            direction: TruncationDirection::default(),
         }
     }
 }
@@ -68,6 +74,12 @@ pub enum TruncationStrategy {
     OnlySecond,
 }
 
+impl Default for TruncationStrategy {
+    fn default() -> Self {
+        TruncationStrategy::LongestFirst
+    }
+}
+
 impl std::convert::AsRef<str> for TruncationStrategy {
     fn as_ref(&self) -> &str {
         match self {
@@ -325,4 +337,13 @@ mod tests {
         truncate_and_assert(get_medium(), get_medium(), &params, 0, 0);
         truncate_and_assert(get_long(), get_long(), &params, 0, 0);
     }
+
+    #[test]
+    fn test_deserialize_defaults() {
+        let old_truncation_params = r#"{"max_length":256,"strategy":"LongestFirst","stride":0}"#;
+
+        let params: TruncationParams = serde_json::from_str(old_truncation_params).unwrap();
+
+        assert_eq!(params.direction, TruncationDirection::Right);
+    }
 }