embarked fix

Artur-Galstyan · Apr 5, 2024 · ecc9096 · ecc9096
1 parent 5f763d3
commit ecc9096
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/jaxonloader/datasets/_datasets.py b/jaxonloader/datasets/_datasets.py
@@ -198,7 +198,19 @@ def _gender_to_int(df: pl.DataFrame) -> pl.DataFrame:
         )
         return df
 
+    def _embarked_to_int(df: pl.DataFrame) -> pl.DataFrame:
+        df = df.with_columns(
+            pl.when(pl.col("Embarked") == "S")
+            .then(0)
+            .when(pl.col("Embarked") == "C")
+            .then(1)
+            .otherwise(2)
+            .alias("Embarked")
+        )
+        return df
+
     train = _gender_to_int(train_df)
+    train = _embarked_to_int(train)
     train_data = train.select(pl.exclude("Survived")).to_numpy()
     train_target = train.select(pl.col("Survived")).to_numpy()
 

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "jaxonloader"
-version = "0.3.2"
+version = "0.3.3"
 description = "A dataloader, but for JAX"
 readme = "README.md"
 requires-python ="~=3.10"