The what

Last year I blogged about img2lego, now renamed brickportraits to remove all connection to LEGO. In my tool I make something that’s not part of the LEGO franchise but can be built by any type of bricks.

Today I got a website brickportraits.londogard.com which allows anyone to run the generation, including a 2D mode (mosaic).

The comparison

So how good is the result? It depends on what you compare too.

Generalist LLM’s (ChatGPT and Gemini) is strong at generating images of 3D Brick Figures, but they hallucinate a lot including the available bricks! This includes the latest Nano Banana 2.

Decent results, but not a fair comparison as it’s not actual bricks nor is it possible to get instructions.

Moving forward I ask competitors to generate an actual 3D Figure brick by brick (ThreeJS). They break down completely, let’s review!

I’m not sure what you think, but to me the winner is as clear as day! 😉

The how

My results are made by AI models built for 3D generation, similar to Meshy.ai, which I voxelize and apply algorithmic enhancements to. This type of model is trained and uses optimal 3D representations internally which makes it easier to work with.

Imagine yourself writing a block by block down and reasoning about how it’ll look, not easy right? It’s the same for the generalist LLM’s.
With that said generalist LLM’s are capable and display good spatial ability in MineBench when generating scenes from text. I’m not sure if the LLM actually reason or is able to generalize from Minecraft content online but for sure each new SotA model accelerate forward!

But! For now this capability is really only available when generating from text, if I add a reference image it fails utterly as shared in my earlier examples.

What do I think personally? Text is not as “rigid” and allows more interpretation which in turn is hallucination-friendly, which can help the model generalize from internet content to produce great results.

Outro

Interested to learn more or discuss something? Please reach out to me!

Dependent Environments

This is not a pixi-only feature, but it’s great to be able to have multiple environments in our base. We use a mono-repo style of working where we combine everything from Training to Monitoring in the same repo.

We set up a base environment which is shared among the sub-environments, this includes dependencies like polars, deltalake and s5cmd which is integral to our way of working.
Then we extend base with computer_vision and add our CV related dependencies, and for our data_pipeline we extend with data utilities.

What’s great about this? We bump polars in one location and it’ll be validated across all environments at the same time.

What’s also great? We can separate dependencies based on platform, i.e. osx-arm64 will have a separate dependency from linux-64 which is important as osx-arm64 doesn’t have CUDA.

Pixi Tasks

Another featuer we’ve utilized more and more with time is tasks. They’re easy yet powerful.

What type of tasks do we have this far?

1. pixi run deploy_monitor: deploy our Streamlit App to Snowflake
2. pixi run labelme: open our labeling program with the correct settings out-of-the-box, and refreshes the pre-signed s3-urls (yes, we’ve patched LabelMe to support URLs :wink:)
3. pixi run docker_deploy: build our docker containrs and push correctly to AWS based on arguments

With one word: Awesome.
We’ve found it helpful in getting people setup and launching internal tools as there’s no need to remember where script is and what order to run.

Temporary Environments

Both Pixi and UV supports temporary environments and global tools. To test something out run:

Final thoughts one year later

Pixi has been a huge success internally in my team, and I’ve been using both pixi and uv privately. Privately I tend to gravitate towards uv as I only have my own environment and it’s a little bit “easier” (hard to explain).

The idea of introducing a completely new tool, pixi, and hearing no complaints, but rather positive remarks is incredible and speaks magnitudes about how good pixi behaves in enterprise settings.
To give some background the team started with broken pip-installs, “works on my computer”, 2 years ago and I transitioned the team into micromamba which stabilized a lot. But it wasn’t a fully positive experience in the team. 1 year ago I moved the team to pixi which has been a lot smoother and happy faces all around.

Happy team, happy life!
~Hampus Londögård

Quick benchmarks:

Please note that this benchmark is not scientific, it’s a simple timeit(number=100), but it’s quite telling anyhow!

Mode	Timer	Person	Chess Board
Original	N/A
PIL.resize(LANCOZ, reducing_gaps=None)	9.18
PIL.resize(LANCOZ, reducing_gaps=2)	0.0003
ski.resize(aa=True)	69.15
cv2.resize(INTER_AREA)	5.88
cv2.resize(LANCOZ + GaussBlur)	4.25
cv2.resize(LANCOZ)	0.033

~Hampus Londögård

How to: Object Detection Inference

infer.py

1from transformers_js_py import import_transformers_js, as_url
transformers = await import_transformers_js()
pipeline = transformers.pipeline

img = "<URL_OR_PATH_TO_AN_IMAGE>"
2pipe = await pipeline("object-detection")

3pred = await pipe(as_url(img))
print(pred) # list of predictions [{"score": float, "label": str, "box": dict[str, int]}, ...]

1

Import library and “import” pipeline

2

Set up pipeline object, downloading model and making things ready to run

3

Run inference. as_url converts local/virtual files to a URL as pipeline object requires URL’s that can be opened in the JS context.

I share how to customize the inference in Section 3 and a full-fledged WebApp with on-device inference using Marimo in Section 4

Why transformer.js.py + pyodide

There’s a good question here: why not run JS directly?
I don’t have a great answer, it’s all about trade-offs.

JS enables “native” usage which likely works better in real-time applications, as it runs JS->WASM rather than WASM->JS->WASM.
What JS doesn’t have is a robust data science ecosystem, unlike Python. “Merging” the two through Pyodide makes sense, and further its fun! 🤓

Inference Customization

To select a specific model define the name as you build the pipeline.

infer_options.py

pipe = await pipeline(
1  "object-detection", # task_name
2  "Xenova/yolos-tiny", # model_name
3  {
    dtype: "q4", # default "q8" for WASM.
    device: "webgpu", # default "WASM" (cpu)
    use_external_data_format: "false" # default "false", set "true" to load >= 2GB model
  }
)

1

Find all available tasks and their linked model-list here.

2

Find all available Object Detection models here.

3

Find all options here and here.

Simple right?
I’m continuously impressed by how far we’ve gotten. On-device inference, even with acceleration, is a painless thing today. If you want simplicity I recommend web and otherwise to use the mobile/native releases or alternatively LiteRT (previously TFLite).

What’s left?
Improving the JS data science ecosystem, for now I prefer Pyodide because of the vast ecosystem. Though I’d like to congratulate transformers.js at successfully making inference simple for people who simply wants a blackbox. Personally I usually want to work with data before/after inference which requires better tools that Pyodide provides.

WASM App using Marimo

If you’ve read my blog you know I recently discovered Marimo, and as always with new tools you try to use them, perhaps a bit too much, whenever you can.
I thought I’d give it a shot to integrate with transformer.js.py and run the inference fully on-device with WASM.
It’s certainly not real-time, but ~5 seconds per image is OK I’d say.

All in all I think this approach is quite neat and could provide very useful, especially for Proof-of-Concepts or Internal Tooling.

Run the app yourself via my marimo.io WASM notebook. Show the code by clicking the three dots in top-right corner.

Thanks for this time,
Hampus Londögård

How to: Object Detection Inference

infer.py

1from transformers_js_py import import_transformers_js, as_url
transformers = await import_transformers_js()
pipeline = transformers.pipeline

img = "<URL_OR_PATH_TO_AN_IMAGE>"
2pipe = await pipeline("object-detection")

3pred = await pipe(as_url(img))
print(pred) # list of predictions [{"score": float, "label": str, "box": dict[str, int]}, ...]

1

Import library and “import” pipeline

2

Set up pipeline object, downloading model and making things ready to run

3

Run inference. as_url converts local/virtual files to a URL as pipeline object requires URL’s that can be opened in the JS context.

I share how to customize the inference in Section 3 and a full-fledged WebApp with on-device inference using Marimo in Section 4

Why transformer.js.py + pyodide

There’s a good question here: why not run JS directly?
I don’t have a great answer, it’s all about trade-offs.

JS enables “native” usage which likely works better in real-time applications, as it runs JS->WASM rather than WASM->JS->WASM.
What JS doesn’t have is a robust data science ecosystem, unlike Python. “Merging” the two through Pyodide makes sense, and further its fun! 🤓

Inference Customization

To select a specific model define the name as you build the pipeline.

infer_options.py

pipe = await pipeline(
1  "object-detection", # task_name
2  "Xenova/yolos-tiny", # model_name
3  {
    dtype: "q4", # default "q8" for WASM.
    device: "webgpu", # default "WASM" (cpu)
    use_external_data_format: "false" # default "false", set "true" to load >= 2GB model
  }
)

1

Find all available tasks and their linked model-list here.

2

Find all available Object Detection models here.

3

Find all options here and here.

Simple right?
I’m continuously impressed by how far we’ve gotten. On-device inference, even with acceleration, is a painless thing today. If you want simplicity I recommend web and otherwise to use the mobile/native releases or alternatively LiteRT (previously TFLite).

What’s left?
Improving the JS data science ecosystem, for now I prefer Pyodide because of the vast ecosystem. Though I’d like to congratulate transformers.js at successfully making inference simple for people who simply wants a blackbox. Personally I usually want to work with data before/after inference which requires better tools that Pyodide provides.

WASM App using Marimo

If you’ve read my blog you know I recently discovered Marimo, and as always with new tools you try to use them, perhaps a bit too much, whenever you can.
I thought I’d give it a shot to integrate with transformer.js.py and run the inference fully on-device with WASM.
It’s certainly not real-time, but ~5 seconds per image is OK I’d say.

All in all I think this approach is quite neat and could provide very useful, especially for Proof-of-Concepts or Internal Tooling.

Run the app yourself via my marimo.io WASM notebook. Show the code by clicking the three dots in top-right corner.

Thanks for this time,
Hampus Londögård

App:Converter

There’s multiple libraries suppoted in Pyodide, Python WASM, among them: polars, duckdb and pandas.
All these libraries are exceptional, with pandas as a exception 😉, to do Data Science and work with tabular data. They also have read/write support for JSON, CSV and parquet files.

These tools enable my simple converter that allows more people to easily move from CSV/JSON to Parquet, and in turn have faster plotting!

As always there’s some problems implementing:

• polars doesn’t support parquet/JSON in WASM
  • –> fall-back to duckdb.
  • duckdb can’t read parquet fromio.BytesIO
    • –> fall-back to… pandas 🤦‍♂️.
    • Luckily we can quickly call pl.from_pandas to run polars!

I hope to add more formats moving forwards, such as ndjson and xlsx.

App:Explore

Data Exploration - a important part and initial step when working with data of any type.
When exploring your dataset it’s good to have a streamlined way of working. Marimo has some excellent tooling to quickly structure your data. I’ve added all these tools in my simple WASM app, things include like:

These tools combine into a quite neat exploration app. If you run this notebook locally you can easily hit up the code cells and modify the DataFrames manually and keep utilizing the nifty UI features such as statistics in DataFrame columns or visualization.

All in all this is a simple quick-starter, I think this app can be helpful for those who wants to explore their data, advaned or simple.

Result

Outro

I’ll keep adding more WASM apps with time. I love Pyodide.

Combining WASM with Marimo, or stlite, feels like such a natural fit.
Marimo to me combines the perfection of Notebook Exploration with Apps, hence I opted for Marimo now.

Moving on I’ll add more in-depth blogs about Marimo and why it’s awesome, embedding WASM snippets and more.

Thanks for this time,
Hampus Londögård

Why I built this

When I was deploying/building a MLFlow Model I found that their infer_code_paths functionality is bugged, as shared in mlflow/issues/14071 and my blog about MLFlow Models, and that I needed something better to really recursively fetch dependencies.

I found that through my nifty little script I could do this better than mlflow themselves. By updating the load_context function we could infer the modules by importing them, assisting mlflow’s infer_code_paths function.

def load_context(self, context):
    # MLFlow bug where parent class is not added to `infer_code_paths`.
    # https://github.com/mlflow/mlflow/issues/14071
    imports = get_dependencies(type(self))
    for module in imports.modules:
        importlib.import_module(module)

This is it for this time,
Hampus Londögård

MLFlow Model

MLFlow Models is MLFlows “self-contained model” that can automatically build a Docker Container and run inference through the built-in mlflow serve command.
It’s an interesting concept that’s not “phenomenal” or “innovating” but helps streamlining our lives, just like the “bread and butter” MLFlow Experiments. I love projects that make the average persons life easier. Advanced user, like I’d call myself, might find it “blocking” but the PythonModel concept I explain later is likely helpful for anyone out there!

MLFlow really hits that sweet point of keeping things simple and not going too far, except perhaps in the current LLM tracing which feels like the shot-gun methodology.

class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input: np.ndarray, params: dict | None):
        # model_input can also be pd.DataFrame, dict[str, np.ndarray], ...
        return model_input

import mlflow

model_path = "my_model.py"

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        python_model=MyModel(),
        artifact_path="my_model",
    )

my_model = mlflow.pyfunc.load_model(model_info.model_uri)

mlflow models build-docker -m runs:/<run_id>/model -n <image_name> --enable-mlserver

import mlflow

mlflow.models.build_docker(
    model_uri=f"runs:/{run_id}/model",
    name="<image_name>",
    enable_mlserver=True,
)

Outro

MLFlow Models provide a simple way to deploy models in a self-contained way.

I hope you try out MLFlow Models as they could end up helping you a lot.

~Hampus Londögård

How do I use Gradio Client?

Step 1: Connect to a Client

from gradio_client import Client, handle_file

client = Client("abidlabs/whisper")

Step 2: Predict

client.predict(
    audio=handle_file("audio_sample.wav")
)

>> "This is a test of the whisper speech recognition model." 

Easy right?!

If we have multiple boxes / steps in the Gradio App we can call each of the components. The client usage of any Gradio App is easily found at the bottom.

And that opens a new menu which shows how to use each box of the App.

Available Clients

Anyone can run this as it supports regular REST requests! They supply a curl sample on how to query the App. But they’ve got a Python and JS native client.

Possibilities

Gradio implemented the whole thing in a way where you don’t need to do anything at all. In turn we now get user-friendly App that is easily deployed with an automatically included REST API. It’s two birds with one stone!

Combining this “App+API” deployment with the free, or paid, HuggingFace Spaces creates a high-value package!

Outro

This blog is really short and to the point. The Gradio Client is sweet and I really wanted to share the experience.
The combination of App+API written at the same time is exciting, when enhancing it with HuggingFace Spaces it all becomes magical. And Gradio’s “auto-share” when running locally that can then provide an API too is not too shabby 😉

This is it for this time,
Hampus Londögård

Usage

I’ll start by sharing my current go-to tool(s) for each area of use, and then try to fit Marimo into this.

• Heavy Applications: Solara
  • Pros: ‘React’ style of programming, very efficient bindings and updates
  • Cons: Not the most modern UI,
• Simple Applications: Streamlit
  • Pros: Simple, Modern UI, Large Community
  • Cons: The “execute everything on each change” execution is quite inefficient and with caching reasoning grows harder with time
• WASM Apps: Streamlit via stlite
• ML Demos with API: Gradio (and Streamlit)
  • Pros: Simple, provides gradio_client REST API by default (AMAZING)
  • Cons: I hate building Gradio apps
• Notebook: Jupyter

Marimo could possibly replace most in the list, but especially WASM Apps and Notebook. Potentially it can take on ML Demos and Simple Applications too, and why not Heavy?

Marimo is perhaps too bold at trying to achieve it all, let’s dive into it!

Quick (Subjective) Rankings

If we put scores on each (3,2,1 for 1st, 2nd, and 3rd) we end up with the following:

1. Marimo & Streamlit: 8pts

2. Solara: 6pts

3. Jupyter & Gradio: 5pts

It seems Marimo ends up covering all needs quite well based on my initial research.
Streamlit ends up in the top because of its strong community, and stlite really helps the WASM story-line.

But how do you actually use Marimo, and can it beat Streamlit by having a smarter execution system?

Marimo

Marimo is easy, you simply define and use a component like following:

# cell 1
import marimo as mo
slider = mo.ui.slider(0,10)
f"Select your step: {slider}"
---
# cell 2
f"You've selected {slider.value} which doubled is {slider.value * 2}"

Here we defined a slider component, we display it using markdown and in our second cell it’s neatly displayed and updated automatically because of reactivity!
Bonus? You can swap the order of the cells and the code will still be valid, because of said reactivity. This is also what enforces the reproducibility. The code follows a DAG based on the variables.
Drawback? You can’t update a variable outside the cell that defines it.

marimo edit --watch 

Outro

I think Marimo all in all does really well, there’s a few sharp edges to resolve but I might replace Jupyter really soon with this. The DataFrame Explorer - Amazing. The callbacks for charts and more - Superb!

It’s like a harder-to-reason but better Streamlit if that makes sense? With more components it’ll be golden!

Thanks for this time,
Hampus Londögård

Result / Demo

I built an Gradio application, mainly as a learning excercise. Usually I prefer Streamlit or Solara for my applications, but Gradio is growing fast and I wanted to try it in a “real project”.

Not to brag, but the results are quite magical in my opinion.
Intrigued? Keep reading to learn how I built this image-to-LEGO generator!

Project Layout

My first step was to research existing solutions. While there are tools for image-to-3D model conversion and voxelization, I couldn’t find anything specifically to build LEGO models with accurate, buildable bricks and “instructions”. This meant I had to build most of the pipeline from scratch. Here’s the steps I identified:

1. Turn image into 3D object (Deep Learning): Use a deep learning model to generate a 3D mesh from a single 2D image.
2. Voxelize Mesh (Algorithmic): Convert the 3D mesh into a voxelized representation, essentially breaking it down into cubes.
3. Colorize Voxels with LEGO approved colors (Algorithmic): Map the colors of the voxels to the closest matching colors from a predefined LEGO color palette.
4. Merge Voxels into LEGO sized bricks (Algorithmic): Combine adjacent voxels into standard LEGO brick shapes, optimizing for larger bricks while preserving the overall shape.
   1. This involves finding the right balance between using larger, more efficient bricks and accurately representing the details of the original model.

This list might make it sound easy, but it really is an interesting and challenging problem!

def voxelize(mesh_path: str | Path, resolution: int):
    """Voxelize the mesh based on a resolution parameter. Resolution is how many bricks it should contain, i.e. 16 creates a 16x16 base plate."""
    mesh = trimesh.load(mesh_path)
    bounds = mesh.bounds
    voxel_size = (bounds[1] - bounds[0]).max() / resolution  # pitch
    voxels = mesh.voxelized(pitch=voxel_size)

def tree_knearest_colors(mesh, voxels):
    tree = cKDTree(mesh.vertices)
    _, vertex_indices = tree.query(voxels.points, k=1)

    return mesh.visual.vertex_colors[vertex_indices]

Future Work

There’s still a lot of work to improve the resulting lego model. The biggest flaws are that there’s too many small bricks and that the colors are not great…

Outro

This was (and is) a very exciting project that I’ll keep on working on whenever I get a short burst of time. There’s some important improvements to add before the project is truly User Friendly™.

The research in what type of approaches are available was a fun one and to implement the algorithms brought me back to my time at Apple, but with the knowledge gained from working in a more mathematical setting later, i.e. applying vectorization.

Finally the result is impressive and something I can use already today. I can’t wait to order bricks for my first LEGO build!

Thanks for this time,
Hampus Londögård

Pixi Docker Builds

After solving your local environment in a easy yet producible manner the next step is to solve it for your cloud workloads - containerization.

Containerization is an important part of a developers toolkit in the modern world. To run cloud workloads it’s very common to deploy as a container, in Data Science this is for everything like Training, Inference and Data Pipelines.

With pixi it’s quite straight-forward and they provide ready-to-use images through the pixi-docker registry. There’s multiple base-images, including CUDA, to get started - it can’t be any simpler!

docker pull ghcr.io/prefix-dev/pixi:latest

Pixi Multi-Platform/Environment

One of Pixi’s standout features is its seamless support for multi-platform and multi-environment projects. While I initially planned to delve deeper into this, prefix.dev recently published an excellent guide on the topic. I highly recommend checking out their documentation on combining different OS’s and environments (CPU, CUDA) with PyTorch for a comprehensive overview.

Pixi Build Slimmming

When containerizing your code, it’s crucial to keep builds slim. Here are a few tricks to help you minimize your Pixi-based Docker images:

1. Leverage .dockerignore: Create a .dockerignore file to exclude unnecessary files and directories (e.g., .git, __pycache__, tests) from your Docker build context.
2. Optimize Dependencies:
   • Carefully consider each dependency and remove any that are not strictly required for production.
   • Utilize multiple environments within your pixi.toml, e.g. prod and dev environments. This allows you to exclude dev-specific dependencies (test, lint, ..) from your production container.
3. Employ Multi-Stage Docker Builds: Multi-stage builds reduces the image size. Use a build stage to install dependencies and compile your application, and then copy only the necessary artifacts to a smaller, leaner final image. The pixi-docker project provides guidance on using multi-stage builds with shell-hook.

Pixi vs uv

While uv has gained significant traction in the Python community, I believe Pixi offers a more compelling solution for my specific needs, especially when it comes to complex, real-world projects.

Why?

1. tasks are awesome. They might not be perfect but they’re great to me!
2. Multi-platform and Multi-environment projects (personal opinion) somehow ends up easier in Pixi
   • I really tried to embrace the uv approach as I appreciate it as more lightweight. But Pixi is somehow “smoother”.
3. Pixi has base-images with CUDA
   • Both tools are easy to build from a raw base-image too, so it’s not a huge problem
4. Access to conda packages
   • Some hate it, but I like getting pre-built binaries.
   • It’s quite interesting to install shell tools via conda for container deployment.
5. Possible to work with other languages than Python

# /// script
# dependencies = [
#   "requests<3",
#   "rich",
# ]
# ///

import requests
from rich.pretty import pprint

resp = requests.get("https://peps.python.org/api/peps.json")
data = resp.json()
pprint([(k, v["title"]) for k, v in data.items()][:10])

Outro

If you’re a Python developer struggling with dependency management, environment inconsistencies, or cumbersome container builds, I urge you to give Pixi a try. It’s a powerful tool that has the potential to streamline your workflow and make you a happier developer. Pixi has certainly made a significant difference in mine!

Thanks for this time, Hampus Londögård

Quick Introduction

Hugging Face Datasets offers easy access to a vast library of datasets, with efficient memory handling through streaming and memory-mapping. Its API simplifies data loading and transformation for direct use with PyTorch and TensorFlow.

Ray Data enables scalable, distributed data processing across multiple nodes, ideal for large datasets. It integrates with Ray’s ML tools for parallel training and distributed transformations. It’s the tool for Large Language Model training, even embraced by OpenAI in their ChatGPT source.

Daft is a high-performance data processing library with lazy evaluation, optimized for structured data formats like Parquet and Arrow. It’s a strong choice for single-node and multi-node data preparation with PyTorch compatibility. It utilizes Ray to achieve multi-node behavior.

PyTorch’s Dataset and DataLoader offer a simple and flexible way to load data with minimal memory overhead, ideal for in-memory and custom datasets. It’s lightweight but lacks distributed and lazy loading features.

Mini Benchmark

Running on full sized images we get a bit more interesting results:

Developer Experience (DX)

def test_elem_by_elem(num_workers: int | None, pin_memory: bool | None, cache: bool):
    if not cache:
        datasets.disable_caching()

    def _preprocess(data: dict):
        imgs = [utils.PREPROCESS_TRANSFORMS(x.convert("RGB")) for x in data["image"]]
        data["image"] = imgs
        return data

    ds = datasets.load_from_disk("./imagenette_full_size")

    def _augment(data: dict):
        tensor = _preprocess(data["image"])
        data["image"] = utils.AUGMENTATIONS(tensor)
        return data

    ds_train = ds["train"].with_transform(_augment)
    ds_valid = ds["validation"].with_transform(_preprocess)

    kwargs = dict(
        num_workers=num_workers or 0,
        persistent_workers=bool(num_workers),
        pin_memory=pin_memory,
        batch_size=32,
    )
    dls_train = torch.utils.data.DataLoader(ds_train, **kwargs)
    dls_valid = torch.utils.data.DataLoader(ds_valid, **kwargs)
    

class ImagenetteDataset(Dataset):
    def __init__(self, hf_dataset, preprocess=None, augment=None):
        self.hf_dataset = hf_dataset
        self.preprocess = preprocess
        self.augment = augment

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        data = self.hf_dataset[idx]

        image = data["image"].convert("RGB")

        # Apply preprocessing and augmentation if specified
        if self.preprocess:
            image: torch.Tensor = self.preprocess(image)
        if self.augment:
            image = self.augment(image)

        return {
            "image": image,
            "label": data["label"],
        }
train_dataset = ImagenetteDataset(
    ds["train"], preprocess=utils.PREPROCESS_TRANSFORMS
)
print(len(train_dataset))
valid_dataset = ImagenetteDataset(
    ds["validation"], preprocess=utils.PREPROCESS_TRANSFORMS
)

# Create DataLoader instances
kwargs = dict(
    num_workers=num_workers or 0,
    persistent_workers=bool(num_workers),
    pin_memory=pin_memory,
    batch_size=32,
)
dls_train = DataLoader(train_dataset, shuffle=True, **kwargs)
dls_valid = DataLoader(valid_dataset, **kwargs)

def load_imagenette_datasets_daft(dataset_path="./imagenette_full_size"):
    ds = datasets.load_from_disk(dataset_path)
    extract_img_bytes = daft.col("image").struct.get("bytes").alias("image")
    ds_train = daft.from_arrow(ds["train"].data.table).select(
        "label", extract_img_bytes
    )

    ds_valid = daft.from_arrow(ds["validation"].data.table).select(
        "label", extract_img_bytes
    )

    img_decode_resize = (
        daft.col("image").image.decode(mode="RGB").image.resize(224, 224)
    )

    ds_train = ds_train.with_column("image", img_decode_resize)
    ds_valid = ds_valid.with_column("image", img_decode_resize)

    def to_f32_tensor(ds: daft.DataFrame):
        return ds.with_column(
            "image",
            daft.col("image").apply(
                lambda x: (x / 255.0).transpose(2, 0, 1),
                return_dtype=daft.DataType.tensor(daft.DataType.float32()),
            ),
        )

    ds_train = to_f32_tensor(ds_train)
    ds_valid = to_f32_tensor(ds_valid)

    ds_train = ds_train.to_torch_iter_dataset()
    ds_valid = ds_valid.to_torch_iter_dataset()

    dls_train = torch.utils.data.DataLoader(ds_train, batch_size=32)
    dls_valid = torch.utils.data.DataLoader(ds_valid, batch_size=32)
    return dls_train, dls_valid

def load_imagenette_datasets_ray(dataset_path="./imagenette_full_size"):
    # Load the Arrow dataset with Ray
    ds = datasets.load_from_disk(dataset_path)

    def extract_img_to_pil(data):
        image = data["image"]["bytes"]
        data["image"] = PIL.Image.open(io.BytesIO(image)).convert("RGB")
        return data

    ds_train = ray.data.from_huggingface(ds["train"]).map(extract_img_to_pil)
    ds_val = ray.data.from_huggingface(ds["validation"]).map(extract_img_to_pil)

    preprocess_transforms = transforms.Compose(
        [
            utils.PREPROCESS_TRANSFORMS,
        ]
    )
    augmentation_transforms = utils.AUGMENTATIONS

    # Apply transformations in Ray
    def preprocess_image(batch):
        batch["image"] = [preprocess_transforms(x) for x in batch["image"]]

        return batch

    def augment_image(elem):
        elem["image"] = augmentation_transforms(elem["image"])

        return batch

    ds_train = ds_train.map_batches(preprocess_image).map(augment_image)

    ds_val = ds_val.map_batches(preprocess_image)

    d_train = ds_train.to_torch(
        label_column="label",
        batch_size=32,
        local_shuffle_buffer_size=512,
        prefetch_batches=5,
    )
    d_valid = ds_val.to_torch(label_column="label", batch_size=32, prefetch_batches=5)

    return d_train, d_valid

I think most of the frameworks ends up at a similar place in the experience.

Quick DX Ranking

1. PyTorch & HuggingFace Datasets
2. Daft
3. Ray (albeit I believe it to be the most scalable solution as you can truly tinker in detail)

I enjoyed Daft a lot with its multi-modal syntax, inspired by polars with namespaces (e.g. .image.decode()), which can be phenomenal. Working with DataFrame’s is a cool addition, where you can drop into python simply by using apply.
Working with Daft more and more I noticed that the DataFrame syntax sometimes becomes a big blocker and the simplicity of HF Datasets and Ray using dict’s in .map statements results in easier code and smoother integration with existing libraries.
Additionally HF Datasets / PyTorch DataLoaders feels more pythonic, where the latter is real simple. I can’t put my finger on it but they just seem easier to debug and understand.

It’ll sure be interesting to follow the progress being made, and I’m happy the dust isn’t settled yet!

Daft

Today I’ll introduce one of the newer alternatives in the field, daft.

Daft is what you can only call a merger between polars, spark and Deep Learning. If they had been more inspired by polars in the Developer Experience (DX) I’d have called it a “lovechild”, but for now they don’t have the nice-to-haves like pl.with_column(new_col_name=pl.col("other_col")*2) named syntax and other things like pl.col("col").replace(dict_to_replace) and a lot of other things.

What daft does have is a multi-modal namespace, unlike polars which solely focuses on traditional data-types. This is really interesting albeit not that fleshed out yet. It’s enjoyable and has potential to grow!

Further, to quote daft themselves:

Daft provides a snappy and delightful local interactive experience, but also seamlessly scales to petabyte-scale distributed workloads.

The petabyte-scale comes from the fact that you can run daft on top of Ray which is a distributed framework that tries to take on Spark. It’s famously used at OpenAI while training their models.

Coding with Daft

Coding with daft is an experience. I only ran locally but it held up really well to “native” PyTorch, even surpassing it in one case!

I’ll share my experience and implementations below!

ds_train = daft.from_arrow(ds["train"].data.table)

df = daft.read_deltalake("some-table-uri") # read_(csv|parquet|json|...)

# expression: lazy non-executed method
extract_img_bytes = daft.col("image").struct.get("bytes").alias("image")

ds_train.select("label", extract_img_bytes)

img_decode_resize = daft.col("image").image.decode(mode="RGB").image.resize(224, 224)

ds_train = ds_train.with_column("image", img_decode_resize)

def rescale_transpose(x: np.array):
    return (x / 255.0).transpose(2, 0, 1)

ds_train.with_column(
    "image",
    daft.col("image").apply(
        rescale_transpose,
        return_dtype=daft.DataType.tensor(daft.DataType.float32()),
    ),
)

ds_train = ds_train.to_torch_iter_dataset()

# NOTE: don't apply num_workers even though PyTorch warns!
dls_train = torch.utils.data.DataLoader(ds_train, batch_size=32)

Mini Benchmark

Comparing speeds with “native” PyTorch DataLoaders is interesting and shows that Daft is on-par in speed when using their new native execution engine (swordfish). When I increase image size, i.e. larger data to process, I see Daft even surpassing PyTorch DataLoaders (!).

N.B. I’m running the full training from a HuggingFace Dataset backed by Arrow. It’s the same underlying data structure for all tests except “Folder File” one, but things might just be different if we start discussing file-loading (rather than from bytes) or even remote data.

df.with_column("image", daft.col("path").url.download())

Remote data

Working with remote data is a common and interesting use-case. I think based on this research that daft has a good chance of performing really well, as the local files also did great.

Final Thoughts

Even if daft has a way to go for Deep Learning training it really holds great promise. If they make the export easier to PyTorch and perhaps add TensorFlow I believe it could grow into a valuable competitor to HuggingFace Datasets et. al.

As Ray is what drives OpenAI’s training I believe Daft stands on some really good scalable underlying tech and can perhaps be what joins Data Engineering and Data Science together as one, for real - a big leap forward!

Thanks for this time, Hampus

Extra: all code is available on the git-repo for this blog, see code/data_loading.

Similarities:

There’s a lot of similarities, it’s quite easy to get started.

UI

One of the more important parts of a tool is the UI. Here I believe in a way ZenML is strong as they “off-load” each components UI to the component itself, i.e. MLFlow tracing is shown in MLFlow UI.

The UI itself of each tool, i.e. WandB, is much better than ClearML’s offering in my opinion.
But the integration of ClearML as a tool “solve all” is a HUGE timesaver and I think could outweigh using the “better” tooling. Integrating everything from Experiment Comparison to Report Building is an quite amazing feat that I think is worthwhile applauding.

Conclusion

First and foremost, I see both Open Source offering moving more and more towards a SaaS. This is clearly visible by locking certain features in the UI (ZenML, the new UI is beautiful but locked down without your Cloud offering). It’s also shown by supplying additional superb features even when self-hosted. I do understand the need to pay your bills, but it’s sad to see Open Source moving to this either way.

See ZenML comparison (Open Source <> Cloud) and ClearML one.

Sometimes the best option is to opt for the “cloud-native” one, i.e. AWS/Azure/GCP tools. But I love open source… :)

Anyhow, to finalize here’s my judgement:

• If you prefer to keep your stack as simple as possible: ClearML.
• If you prefer to keep your stack customized having the best tool for each part: ZenML.

I cannot pick a winner, ZenML enables simpler transition and better tooling all in all, but the full-on integration of ClearML with “everything working together” is quite magical and similar to the cloud-native options (AWS Sagemaker/Azure MLStudio/GCP Vertex).

Find the code for each framework running MNIST: ….

Thanks for this time, Hampus

task = Task.init(
    project_name="MNIST Digit Recognition",
    task_name="Simple NN model with PyTorch Lightning",
    task_type=Task.TaskTypes.training,
    output_uri=None,
)


class SimpleNN(pl.LightningModule):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(784, 512)
        self.dropout = nn.Dropout(0.2)
        self.fc2 = nn.Linear(512, 10)

    def forward(self, x):
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = torch.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return torch.log_softmax(x, dim=1)

    def training_step(self, batch, batch_idx):
        data, target = batch
        output = self(data)
        loss = nn.functional.cross_entropy(output, target)
        self.log("train_loss", loss)
        return loss

    def test_step(self, batch, batch_idx):
        return self(batch[0])

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=0.001)


params_dictionary = {"epochs": 3}
task.connect(params_dictionary)

transform = transforms.Compose(
    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
)

train_dataset = datasets.MNIST(
    root="./data", train=True, transform=transform, download=True
)
test_dataset = datasets.MNIST(root="./data", train=False, transform=transform)

train_loader = torch.utils.data.DataLoader(
    dataset=train_dataset, batch_size=128, shuffle=True
)
test_loader = torch.utils.data.DataLoader(
    dataset=test_dataset, batch_size=128, shuffle=False
)

model = SimpleNN()
trainer = pl.Trainer(max_epochs=params_dictionary["epochs"])
trainer.fit(model, train_loader)
trainer.test(dataloaders=test_loader)

@zenml.step
def load_mnist() -> Tuple[
    Annotated[torch.utils.data.DataLoader, "train_loader"],
    Annotated[torch.utils.data.DataLoader, "test_loader"],
]:
    ...
    return train_loader, test_loader

@zenml.step
def train_model(
    train_loader: torch.utils.data.DataLoader, test_loader: torch.utils.data.DataLoader
):
    model = SimpleNN()
    trainer = pl.Trainer()
    trainer.fit(model, train_loader)
    trainer.test(dataloaders=test_loader)


@zenml.pipeline
def train_pipeline():
    train_loader, test_loader = load_mnist()
    train_model(train_loader, test_loader)


if __name__ == "__main__":
    train_pipeline()

Play Around

First I play around with fragments, testing the most simple use-case – and I’m sold!

N.B. this is already possible in other tools such as Solara that has a better reactive approach, but streamlit has a bigger user-base and I love to see a solution to this long-standing problem!

import numpy as np
import streamlit as st


def main():
    st.write("# Main Function")
    st.write("Hello, World! (main)")
    st.toggle("Toggle me!")


@st.experimental_fragment()
def first_fragment():
    st.write("## First Fragment")
    random_choice = np.random.choice(["a", "b", "c"])
    st.write(f"Random choice: {random_choice}")
    st.toggle("Toggle me! (1st Fragment)")


@st.experimental_fragment(run_every="2s")
def second_fragment():
    st.write("## Second Fragment")
    st.write("Hello, World! (2nd Fragment)")
    random_choice = np.random.choice(["a", "b", "c"])
    st.write(f"Random choice: {random_choice}")
    st.toggle("Toggle me! (2nd Fragment)")


if __name__ == "__main__":
    main()
    c1, c2 = st.columns(2)

    with c1:
        first_fragment()
    with c2:
        second_fragment()

This enables the following behavior:

1. Toggling “main” will refresh everything
2. Toggling a fragment will only refresh that fragment
3. Second fragment will refresh every 2 seconds

What is refreshed? The Random choice letter is updated to a random letter (a, b, or c).

All in all this is what we’d probably do in a Dashboard. See the following GIF’s:

Adding Complexity

As always it’s a lot more fun to test these things in scenarios that are closer to real-life, and that’s what I intend to do!

1. Fetching data from a data storage
2. Displaying different graphs
3. Sharing state from main

In this graph we have a Amplitude Multiplier (main) that affects both fragments, additionally we have a sine wave where the frequency is editable and will only re-render (re-compute) that fragment (first). Finally there’s a Stock Fragment (second) which automatically updates every 2 seconds, unless locked it’ll randomly select a stock, if locked we can still change stock and it’ll only re-render that fragment (second).

See the GIF below! 👇

import numpy as np
import streamlit as st
import polars as pl
import plotly.express as px


def main() -> float:
    st.write("# Main Function")
    st.write("Hello, World! (main)")
    multiplier = st.slider("Amplitude Multiplier", 0.0, 10.0, 1.0, 0.1)

    return multiplier


@st.cache_resource
def get_stocks() -> pl.DataFrame:
    return pl.read_csv(
        "https://raw.githubusercontent.com/vega/datalib/master/test/data/stocks.csv"
    )


@st.experimental_fragment()
def first_fragment(multiplier: float):
    st.write("## First Fragment")
    sine_frequency = st.slider("Sine Frequency", 0.0, 10.0, 1.0, 0.1)
    # create sine wave with multiplier height and sine_frequency as frequency
    t = np.linspace(0, 2 * np.pi * sine_frequency, 100)
    y = multiplier * np.sin(t)

    df = pl.DataFrame({"t": t, "y": y})
    st.plotly_chart(
        px.line(df, x="t", y="y", title="Sine wave"), use_container_width=True
    )


@st.experimental_fragment(run_every="2s")
def second_fragment(multiplier: float):
    st.write("## Second Fragment")
    c1, c2 = st.columns(2)

    with c1:
        if not st.checkbox("Lock company"):
            st.session_state["ticker_select"] = np.random.choice(
                ["AAPL", "GOOG", "AMZN"]
            )
    with c2:
        ticker = st.selectbox(
            "Company (symbol)", ["AAPL", "GOOG", "AMZN"], key="ticker_select"
        )
    stocks = get_stocks()
    stocks = stocks.filter(pl.col("symbol") == ticker).with_columns(
        pl.col("price") * multiplier
    )

    st.plotly_chart(
        px.line(stocks, x="date", y="price", title=f"Stock price ({ticker})"),
        use_container_width=True,
    )


if __name__ == "__main__":
    multiplier = main()
    c1, c2 = st.columns(2)

    with c1:
        first_fragment(multiplier)
    with c2:
        second_fragment(multiplier)

Drawbacks

This solution doesn’t fit every scenario, and as usual with Streamlit, integrating it introduces complexity via state management. Fragments add another level atop the existing st.state, potentially introducing more intricacies and headaches.

Other solutions such as Solara and Panel has this more built into the solution, but then again their entry threshold is a lot higher!

Outro

Any other questions? Please go ahead and ask!

This development is exciting and will for sure give Streamlit new life in “efficiency”. I, for one, am happy to see all new Data Apps fighting!

Finally, all the code is available on this blogs github under code_snippets.

/ Hampus Londögård

Simple get-started

pixi init # create toml and lock files
pixi add python polars # add python and polars as dependencies

pixi shell # activates the virtual environment
# Alternatively, "pixi run python ..."

Add Task

Tasks are really awesome to reduce the threshold to enter projects. It’s simpler than a spread of bash-scripts or other things.

One standardized way to do things! :)

pixi task add gui "solara run app.py" # Adds task

pixi run gui # runs Solara App

Outro

Please read the docs to learn more!