Using LLM for Machine Learning: A Step-by-Step Guide

#learnai #oxlo #ai

We are going to build an LLM agent that turns a plain-English dataset description into a runnable scikit-learn training script. It writes the code, reviews itself for data leakage, and executes the pipeline. If you prototype ML models frequently, this removes boilerplate so you can focus on feature engineering.

What you'll need

Python 3.10 or newer, the OpenAI SDK, and a free Oxlo.ai API key from https://portal.oxlo.ai. Install the dependencies with pip.

pip install openai pandas scikit-learn

Because Oxlo.ai charges a flat rate per request rather than per token, iterating with long system prompts and large code outputs stays predictable. You can explore the details on the Oxlo.ai pricing page.

Step 1: Connect and test the Oxlo.ai client

I start every project by verifying the endpoint and credentials. This snippet initializes the OpenAI-compatible client against Oxlo.ai and pings it with a lightweight prompt.

from openai import OpenAI

client = OpenAI(
    base_url="https://api.oxlo.ai/v1",
    api_key="YOUR_OXLO_API_KEY"
)

response = client.chat.completions.create(
    model="llama-3.3-70b",
    messages=[
        {"role": "user", "content": "Reply with exactly: Oxlo.ai client is ready."}
    ],
)

print(response.choices[0].message.content)

Step 2: Craft the system prompt for ML code generation

The system prompt is the only training the agent gets. I keep it strict: one file, no markdown fences, and mandatory pipeline best practices.

SYSTEM_PROMPT = """You are a senior ML engineer. When given a dataset description, output a single, self-contained Python script that:

1. Loads the CSV with pandas.
2. Separates features and the target column.
3. Uses ColumnTransformer with StandardScaler for numeric columns and OneHotEncoder(handle_unknown='ignore') for categoricals.
4. Splits with train_test_split. Use stratify=y for classification tasks.
5. Trains a model. Use RandomForestClassifier for classification or RandomForestRegressor for regression unless the user specifies otherwise.
6. Evaluates with appropriate metrics and prints them.
7. Saves the trained pipeline to disk as pipeline.pkl.

Do not write markdown code fences. Output only valid Python code."""

Step 3: Build the generator function

This function packages the user request and fires it to Oxlo.ai. I use Llama 3.3 70B because it produces clean, structured Python with minimal hallucinated imports.

def generate_training_script(csv_path, target_col, task_type="classification"):
    user_message = f"""Dataset: {csv_path}
Target column: {target_col}
Task type: {task_type}

Write the full training script according to the system instructions.
"""
    response = client.chat.completions.create(
        model="llama-3.3-70b",
        messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message},
        ],
    )
    return response.choices[0].message.content.strip()

raw_script = generate_training_script("data/iris.csv", "species", "classification")
print(raw_script[:500])

Step 4: Add a review loop to catch data leakage

Generated code can place preprocessing outside cross-validation or leak test data into feature engineering. I pass the draft back to the model with a critique prompt so it fixes its own mistakes.

def review_script(draft):
    critique = f"""Review this sklearn script for data leakage, incorrect train-test split ordering, and metric choice.
If you find issues, return a corrected full script. If it is correct, return it unchanged.
Do not write markdown fences. Output only valid Python code.

Script:
{draft}
"""
    response = client.chat.completions.create(
        model="llama-3.3-70b",
        messages=[
            {"role": "system", "content": "You are a senior ML engineer who audits code for correctness."},
            {"role": "user", "content": critique},
        ],
    )
    return response.choices[0].message.content.strip()

final_script = review_script(raw_script)
print(final_script[:500])

Step 5: Persist the artifact and run it locally

Once the code passes review, I write it to disk and execute it in a subprocess. This keeps the generated work reproducible and version-controllable.

import subprocess

script_path = "train_model.py"
with open(script_path, "w") as f:
    f.write(final_script)

print(f"Saved to {script_path}")

result = subprocess.run(["python", script_path], capture_output=True, text=True)
print(result.stdout)
if result.returncode != 0:
    print("STDERR:", result.stderr)

Run it

Here is a complete invocation that creates a tiny CSV, calls the agent, and executes the result.

import pandas as pd
from sklearn.datasets import load_iris
import subprocess

# Prepare a sample dataset
iris = load_iris(as_frame=True)
df = iris.frame
df["species"] = df["target"].map({0: "setosa", 1: "versicolor", 2: "virginica"})
df.drop(columns=["target"], inplace=True)
df.to_csv("iris.csv", index=False)

# Generate, review, and execute
raw = generate_training_script("iris.csv", "species", "classification")
clean = review_script(raw)

with open("train_iris.py", "w") as f:
    f.write(clean)

subprocess.run(["python", "train_iris.py"])

Typical output looks like this:

Saved to train_iris.py
Train accuracy: 0.97
Test accuracy: 0.95
Pipeline saved to pipeline.pkl

Next steps

Swap llama-3.3-70b for deepseek-v3.2 if you need heavier reasoning on messy tabular data, or try qwen-3-32b for multilingual dataset descriptions. You can also extend the agent by storing generated feature schemas in an Oxlo.ai embeddings index and retrieving them for future pipelines.

DEV Community