diff --git a/functime/base/model.py b/functime/base/model.py
index edc4388d..3b0f697f 100644
--- a/functime/base/model.py
+++ b/functime/base/model.py
@@ -15,7 +15,7 @@ def _set_string_cache(df: pl.DataFrame):
         # Reset categorical to string type
         df = df.with_columns(pl.col(entity_col).cast(pl.Utf8))
     df_new = df.with_columns(
-        pl.col(entity_col).map_dict(string_cache, return_dtype=pl.Int32)
+        pl.col(entity_col).replace(string_cache, return_dtype=pl.Int32, default=None)
     )
     inv_string_cache = {i: entity for entity, i in string_cache.items()}
     return df_new, entity_col_dtype, string_cache, inv_string_cache
@@ -29,7 +29,7 @@ def _enforce_string_cache(
         # Reset categorical to string type
         df = df.with_columns(pl.col(entity_col).cast(pl.Utf8))
     return df.with_columns(
-        pl.col(entity_col).map_dict(string_cache, return_dtype=pl.Int32)
+        pl.col(entity_col).replace(string_cache, return_dtype=pl.Int32, default=None)
     )
 
 
@@ -37,7 +37,9 @@ def _reset_string_cache(
     df: pl.DataFrame, inv_string_cache: Mapping[int, Union[int, str]], return_dtype
 ) -> pl.DataFrame:
     return df.with_columns(
-        pl.col(df.columns[0]).map_dict(inv_string_cache, return_dtype=return_dtype)
+        pl.col(df.columns[0]).replace(
+            inv_string_cache, return_dtype=return_dtype, default=None
+        )
     )