DEV Community: Vaibhav Dubey

[Boost]

Vaibhav Dubey — Mon, 12 May 2025 16:39:15 +0000

Vaibhav Dubey

May 12 '25

Plexe - ML models from a prompt

#ai #opensource #machinelearning #data

Comments

3 min read

Plexe - ML models from a prompt

Vaibhav Dubey — Mon, 12 May 2025 16:34:36 +0000

We’re building Plexe, an open-source agent that turns natural language task descriptions into trained ML models. Here’s a video walkthrough: https://www.youtube.com/watch?v=bUwCSglhcXY.

There are all kinds of uses for ML models that never get realized because the process of making them is messy and convoluted. You can spend months trying to find the data, clean it, experiment with models and deploy to production, only to find out that your project has been binned for taking so long. There are many tools for “automating” ML, but it still takes teams of ML experts to actually productionize something of value. And we can’t keep throwing LLMs at every ML problem. Why use a generic 10B parameter language model, if a logistic regression trained on your data could do the job better?

Our light-bulb moment was that we could use LLMs to generate task-specific ML models that would be trained on one’s own data. Thanks to the emergent reasoning ability of LLMs, it is now possible to create an agentic system that might automate most of the ML lifecycle.

A couple of months ago, we started developing a Python library that would let you define ML models on structured data using a description of the expected behaviour. Our initial implementation arranged potential solutions into a graph, using LLMs to write plans, implement them as code, and run the resulting training script. Using simple search algorithms, the system traversed the solution space to identify and package the best model.

However, we ran into several limitations, as the algorithm proved brittle under edge cases, and we kept having to put patches for every minor issue in the training process. We decided to rethink the approach, throw everything out, and rebuild the tool using an agentic approach prioritising generality and flexibility. What started as a single ML engineering agent turned into an agentic ML "team", with all experiments tracked and logged using MLFlow.

Our current implementation uses the smolagents library to define an agent hierarchy. We mapped the functionality of our previous implementation to a set of specialized agents, such as an “ML scientist” that proposes solution plans, and so on. Each agent has specialized tools, instructions, and prompt templates. To facilitate cross-agent communication, we implemented a shared memory that enables objects (datasets, code snippets, etc) to be passed across agents indirectly by referencing keys in a registry. You can find a detailed write-up on how it works here.
Plexe’s early release is focused on predictive problems over structured data, and can be used to build models such as forecasting player injury risk in high-intensity sports, product recommendations for an ecommerce marketplace, or predicting technical indicators for algorithmic trading. Here are some examples to get you started!

To get it working on your data, you can dump any CSV, parquet, etc and Plexe uses what it needs from your dataset to figure out what features it should use. In the open-source tool, it only supports adding files right now but in our platform version, we'll have support for integrating with Postgres where it pulls all available data based on an SQL query and dumps it into a parquet file for the agent to build models.

Next up, we’ll be tackling more of the ML project lifecycle: we’re currently working on adding a “feature engineering agent” that focuses on the complex data transformations that are often required for data to be ready for model training. If you're interested, check Plexe out and let us know your thoughts!

I built an open-source library to generate specialised machine learning models from natural language

Vaibhav Dubey — Thu, 13 Feb 2025 09:22:14 +0000

Smolmodels is an open-source library that uses graph search with LLM-based code generation to automatically create lightweight, task-specific ML models from natural language descriptions.

Here's how it works with a time-series prediction example; let’s say df is a dataframe containing the “air passengers” dataset from statsmodels.

    import smolmodels as sm

    model = sm.Model(
        intent="Predict the number of international air passengers (in thousands) in a given month, based on historical time series data.",
        input_schema={"Month": str},
        output_schema={"Passengers": int}
    )

    model.build(dataset=df, provider="openai/gpt-4o")

    prediction = model.predict({"Month": "2019-01"})

    sm.models.save_model(model, "air_passengers")

Under the hood, the library:

Parses the intent to identify ML task type and constraints
Uses graph search to explore potential model architectures
Automatically optimises the solutions produced
Generates task-specific training code
Generates inference code to use the model

The project is in alpha stage and completely open source: https://github.com/plexe-ai/smolmodels

The library is fully open-source (Apache-2.0), so feel free to use it however you like. Or just tear it apart in the comments if you think this is dumb. We’d love some feedback, and we’re very open to code contributions!

Created my first open-source project

Vaibhav Dubey — Thu, 30 Jan 2025 09:08:51 +0000

Took about a month to get here but happy to release this open source project I've been working on for the last month: https://github.com/plexe-ai/smolmodels

It lets you automatically build ML models by defining what you need in plain words and setting up input/output schemas. Under the hood it uses graph search and LLMs to explore different models and compare their performance

This is an early alpha, actively being developed. Would love feedback or ideas on where to take it next.