diff --git a/Cargo.toml b/Cargo.toml
index 7956781..3102e14 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -2,7 +2,7 @@
 members = ["python"]
 
 [workspace.package]
-version = "0.1.2"
+version = "0.2.0"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 description = "Token healing implementation"
@@ -25,23 +25,23 @@ readme = "README.md"
 exclude = ["release-plz.toml", ".github", "python"]
 
 [dependencies]
-derive_more = "0.99.17"
+derive_more = { version = "1.0.0", features = ["deref", "as_ref"] }
 general-sam = { version = "1.0.0", features = ["trie"] }
-pyo3 = { version = "0.21.2", optional = true }
+pyo3 = { version = "0.22.5", optional = true }
 smallvec = "1.13.2"
-thiserror = "1.0.60"
+thiserror = "1.0.64"
 
 [features]
 pyo3 = ["dep:pyo3"]
 
 [dev-dependencies]
-clap = { version = "4.5.4", features = ["derive", "env"] }
+clap = { version = "4.5.20", features = ["derive", "env"] }
 color-eyre = "0.6.3"
 rand = "0.8.5"
-regex = "1.10.4"
-serde_json = "1.0.117"
-tokenizers = { version = "0.19.1", features = ["hf-hub", "http"] }
-tokio = { version = "1.37.0", features = ["rt-multi-thread"] }
+regex = "1.11.0"
+serde_json = "1.0.129"
+tokenizers = { version = "0.20.1", features = ["hf-hub", "http"] }
+tokio = { version = "1.40.0", features = ["rt-multi-thread"] }
 
 [package.metadata.docs.rs]
 all-features = true
diff --git a/python/Cargo.toml b/python/Cargo.toml
index 171f443..4df16eb 100644
--- a/python/Cargo.toml
+++ b/python/Cargo.toml
@@ -15,5 +15,5 @@ name = "mtc_token_healing"
 crate-type = ["cdylib"]
 
 [dependencies]
-mtc-token-healing = { version = "0.1.2", path = "..", features = ["pyo3"] }
-pyo3 = { version = "0.21.2", features = ["extension-module", "generate-import-lib", "abi3-py39"] }
+mtc-token-healing = { version = "0.2", path = "..", features = ["pyo3"] }
+pyo3 = { version = "0.22.5", features = ["extension-module", "generate-import-lib", "abi3-py39"] }
diff --git a/src/search_tree.rs b/src/search_tree.rs
index eedc03d..4a9ca3d 100644
--- a/src/search_tree.rs
+++ b/src/search_tree.rs
@@ -340,6 +340,7 @@ mod _pyo3 {
     #[pymethods]
     impl InferResponse {
         #[new]
+        #[pyo3(signature = (sampled=None, sparse_choices=None))]
         pub fn new_py(
             sampled: Option<Prediction>,
             sparse_choices: Option<Vec<Prediction>>,