DEV Community

Cover image for Build a Titanic Model in Python, Then Edit the Pipeline Visually
Edwardvaneechoud
Edwardvaneechoud

Posted on AI-assisted

Build a Titanic Model in Python, Then Edit the Pipeline Visually

#python #polars #datascience #machinelearning

Use Flowfile's Polars-style Python API to build a Titanic model, inspect the pipeline as a graph, and keep editing it in a visual notebook.

I like writing data transformations in Python, but when a pipeline branches into multiple steps, I sometimes miss seeing how everything connects. Furthermore, when I then have to present my changes to my non-technical colleagues I end up drawing the flow in draw.io.

That's something I've been working on with Flowfile, my open-source data pipeline tool. You can write Polars-style Python code, inspect the resulting graph, and keep editing it visually.

In this example, I'm using Kaggle's Titanic competition. We'll build a model, generate predictions, and then open the same pipeline in Flowfile's editor.

Getting started

Flowfile has a visual editor, but we won't start there. We'll build the pipeline in Python using ff.col, with_columns, filter, and other operations that should look familiar if you've used Polars.

Install Flowfile:

pip install "flowfile>=0.23.2"
Enter fullscreen mode Exit fullscreen mode

The Titanic competition gives us two files:

  • train.csv has 891 passengers and a Survived column.
  • test.csv has 418 passengers. We need to predict Survived for those rows.

You can download them from Kaggle, after accepting the competition rules, or use the Kaggle CLI:

kaggle competitions download -c titanic
Enter fullscreen mode Exit fullscreen mode

Extract the files before running the script. The paths below point to my local downloads folder; change them to match where you saved yours.

Build the pipeline

Read and combine the data

We'll start by reading both files and putting them into one frame. That lets us use the same feature transformations for the training and test passengers.

import flowfile as ff

train = ff.read_csv("~\\downloads\\titanic\\train.csv")
test = ff.read_csv("~\\downloads\\titanic\\test.csv")

full = ff.concat([train, test], how="diagonal_relaxed")
Enter fullscreen mode Exit fullscreen mode

test.csv doesn't contain Survived. With diagonal_relaxed, those rows get a null in that column, which we'll use later to separate them again.

Get more information out of the existing columns

Passenger names contain titles such as Mr, Mrs, Miss, and Master, which tell us something about age and social position. We'll extract those, calculate family size, and add two flags.

full = full.with_columns(
    ff.col("Name").str.extract(r",\s*([^.]+)\.", 1).alias("Title"),
    (ff.col("SibSp") + ff.col("Parch") + 1).alias("FamilySize"),
    ff.col("Cabin").is_not_null().cast(ff.Int8).alias("HasCabin"),
    (ff.col("Sex") == "female").cast(ff.Int8).alias("IsFemale"),
)
Enter fullscreen mode Exit fullscreen mode

Some titles appear only once or twice, so we'll keep the four common ones and group the rest under Rare.

full = full.with_columns(
    ff.when(ff.col("Title").is_in(["Mr", "Mrs", "Miss", "Master"]))
    .then(ff.col("Title"))
    .otherwise(ff.lit("Rare"))
    .alias("Title"),
)
Enter fullscreen mode Exit fullscreen mode

Add group context and fill missing values

For missing ages, a title-based median is more useful than one median for everyone: Master and Mr refer to very different age groups. We'll use window expressions for that, and fill fares based on passenger class:

full = full.with_columns(ff.col("Age").median().over("Title").alias("TitleAge"))
full = full.with_columns(ff.col("Fare").median().over("Pclass").alias("ClassFare"))
Enter fullscreen mode Exit fullscreen mode

Now fill the missing ages and fares from the corresponding group. There are also two missing embarkation ports; we'll use S (Southampton) for those.

full = full.with_columns(
    ff.col("Age").fill_null(ff.col("TitleAge")),
    ff.col("Fare").fill_null(ff.col("ClassFare")),
    ff.col("Embarked").fill_null("S"),
)
Enter fullscreen mode Exit fullscreen mode

Why the code looks like Polars

If you've used Polars, most of this should look familiar. ff.col(...) follows the same expression style as pl.col(...), and methods like with_columns, over, and fill_null work in much the same way. The difference is that Flowfile also keeps track of these transformations as a graph. We'll see how they show up as visual nodes later.

Prepare the model features

We'll use logistic regression, so the categorical features need to become numeric columns. I'm keeping this fairly simple: a few 0/1 flags for class, title, and embarkation port.

full = full.with_columns(
    (ff.col("Pclass") == 1).cast(ff.Int8).alias("FirstClass"),
    (ff.col("Pclass") == 2).cast(ff.Int8).alias("SecondClass"),
    (ff.col("FamilySize") == 1).cast(ff.Int8).alias("IsAlone"),
    *[(ff.col("Title") == t).cast(ff.Int8).alias(f"Title{t}") for t in ["Master", "Miss", "Mrs", "Rare"]],
    *[(ff.col("Embarked") == e).cast(ff.Int8).alias(f"From{e}") for e in "CQ"],
)

FEATURES = ["IsFemale", "FirstClass", "SecondClass", "Age", "Fare", "FamilySize", "IsAlone", "HasCabin",
            "TitleMaster", "TitleMiss", "TitleMrs", "TitleRare", "FromC", "FromQ"]
Enter fullscreen mode Exit fullscreen mode

Train and check the model

Because the training and test data are in one frame, we can split them back out using Survived. Then we'll hold out 20% of the labeled rows to get an initial idea of how well the model works.

For the data science part, Flowfile uses polars-ds under the hood, including the logistic regression we're training here.

labeled = full.filter(ff.col("Survived").is_not_null())
unlabeled = full.filter(ff.col("Survived").is_null())

fit, holdout = labeled.random_split({"fit": 80, "holdout": 20}, seed=42)
check = fit.train_model(target="Survived", features=FEATURES, model_type="logistic_regression")
metrics = (
    holdout.wait_for(check)
    .apply_model(upstream=check, output_column="prediction")
    .evaluate_model("Survived", predicted_column="prediction", upstream=check)
)
Enter fullscreen mode Exit fullscreen mode

train_model fits the model when it's called. wait_for makes the dependency explicit: the holdout predictions need the fitted model first.

For the Kaggle submission, we'll train another model on all labeled rows and use it to predict the passengers in test.csv.

model = labeled.train_model(target="Survived", features=FEATURES, model_type="logistic_regression")

# model.read_model() gives you access to the actual model and the underlying python object.

submission = (
    unlabeled.wait_for(model)
    .apply_model(upstream=model, output_column="prediction")
    .select("PassengerId", ff.col("prediction").alias("Survived"))
)
Enter fullscreen mode Exit fullscreen mode

So far we've written the whole pipeline in Python. Before writing out the submission, let's look at the graph and the validation results.

What did we build?

Print the graph

Every operation above added something to Flowfile's graph. We can inspect it without opening the editor:

submission.flow_graph.print_tree()
Enter fullscreen mode Exit fullscreen mode 
● Read data (3)
│ ● Read data (4)
├─┘
● Union data (5)
│
● Polars code (6)  Add columns: Title, FamilySize, HasCabin (+1 more)
│
● Formula (8)
│
● Window functions (9)
│
● Window functions (10)
│
● Formula (12)
│
● Formula (14)
├─┐
● │ Filter data (15)
├─┼─┐
│ ● │ Filter data (16)
● │ │ Random Split (17)
├─┼─┼─┐
● │ │ │ Train Model (18)  [fit]
├─┼─┼─┘
● │ │ Wait For (19)  [holdout]
│ │ │
● │ │ Apply Model (20)
│ │ │
● │ │ Evaluate Model (21)
  │ ● Train Model (22)
  ├─┘
  ● Wait For (23)
  │
  ● Apply Model (24)
  │
  ● Select data (25)
Enter fullscreen mode Exit fullscreen mode

Both inputs merge, the transformations happen once, and the pipeline branches into validation and final predictions.

The transformations also become different kinds of nodes: group medians become Window functions, and conditional expressions and null replacements become Formula nodes. Since str.extract has no matching Formula operation, Flowfile keeps it as a Polars code node. I like that I don't have to rewrite Python just to make the graph editable.

Check the result 

print(metrics.collect())
Enter fullscreen mode Exit fullscreen mode 
┌───────────┬──────────┐
│ metric    ┆ value    │
│ ---       ┆ ---      │
│ str       ┆ f64      │
╞═══════════╪══════════╡
│ accuracy  ┆ 0.842697 │
│ precision ┆ 0.828328 │
│ recall    ┆ 0.832648 │
│ f1        ┆ 0.830361 │
│ n_correct ┆ 150.0    │
│ n_total   ┆ 178.0    │
└───────────┴──────────┘
Enter fullscreen mode Exit fullscreen mode

That's 150 correct predictions out of 178, or about 84% accuracy.

One caveat: we calculated age and fare medians before the split, so the holdout rows influenced preprocessing. The model didn't see their labels, but this isn't fully isolated validation. For that, we'd calculate medians using only the training split. With just 178 holdout rows, the Kaggle score may also differ.

Write the submission

Now that we've inspected the output, let's write out the predictions. submission contains the PassengerId and predicted Survived columns for the rows from test.csv.

submission.write_csv("submission.csv")
Enter fullscreen mode Exit fullscreen mode

That creates submission.csv, which you can upload to Kaggle.

Open the same pipeline in the editor

So far we've stayed in Python. Let's open the graph:

ff.open_graph_in_editor(submission.flow_graph)
Enter fullscreen mode Exit fullscreen mode

Flowfile starts the server if needed and opens the pipeline in your browser. The nodes from print_tree() are now on a canvas, where you can inspect their settings and preview the data at each step.

Overview of the flow

For example, open the Formula node that fills the missing values. You'll see the three expressions we wrote for age, fare, and embarkation port. The graph isn't a separate diagram I made to explain the Python. It's the pipeline itself.

Make a change from the notebook

The code panel has FlowFrame, Polars, project and Notebook tabs. Open Notebook to see the pipeline broken into cells, one per step. You can edit a cell and run it, and the canvas updates to match. The sync is also a single undo step, so you can go back if you change your mind.

Canvas and notebook in one overview

Let's try a feature that isn't in the original script. We already use age as a number, but what if we explicitly mark children? Add this expression inside the feature-engineering with_columns call:

(ff.col("Age") < 12).cast(ff.Int8).alias("IsChild")
Enter fullscreen mode Exit fullscreen mode

Then add "IsChild" to the holdout model's features list and rerun the model and evaluation steps.

The change appears on the canvas, and the model retrains. In my run, accuracy stayed at 0.843. TitleMaster may already capture much of the signal for young boys, so the new flag didn't improve this holdout result. Still, it was easy to try without rebuilding the pipeline in another interface.

Go back to a Python script

When you're done in the notebook editor, open the FlowFrame tab. It gives you the full pipeline as a standalone Python script, including the changes made on the canvas or in the notebook. You can copy that back into your project and run it there.

There's also a Polars tab for exporting plain Polars code. That works for flows made up of operations Polars can represent. This one has Flowfile-specific model training and prediction nodes, so FlowFrame is the export we want here.

Wrapping up

We built the Titanic model in Python, inspected its graph, tried changing a feature in the visual notebook, and exported the result back to code.

That's what I've been trying to make possible with Flowfile: use Python where it makes sense, switch to the canvas when that's easier, and keep working on the same pipeline.

If you want to try it, Flowfile is open source:

I'd be interested to hear where switching between code and a canvas helps on real pipelines, and where it gets in the way.

Top comments (0)