het-ai 1.0.0

Key Features

• Framework-agnostic — optional dependency groups bring first-class support for PyTorch, TensorFlow, scikit-learn, PyMC, and arbitrary black-box processes.

• Broad task coverage — supports high-frequency AI model development workflows across classification, regression, clustering, anomaly detection, time-series forecasting, text classification, Bayesian inference, and black-box optimisation.

• Declarative search spaces — annotate tunable parameters directly in the train() signature with @BaseTrainer.search(...). No Optuna boilerplate required.

• Single- and multi-objective HPO — native Pareto-front support for multi-metric optimisation; override select_best_trial() for custom selection logic.

• Environment-variable-driven config — every TrainConfig field can be overridden by an environment variable, making container and CI/CD integration seamless.

• Dry-run validation — verify the entire pipeline locally without production data or a remote training server using a single trainer.dry_run() call.

• DVC + SeaweedFS integration — het_ai.dvc resolves the latest data version tag from GitHub, pulls actual data from SeaweedFS via DVC, and injects version metadata into the DataBundle for downstream traceability.

• MLflow auto-logging — set TrainConfig(mlflow=MLflowConfig(...)) and the framework automatically logs params, metrics, dataset lineage, and registers the model after every run. No MLflow code required in your Trainer.

• Flexible model export — ONNX / TFLite / Keras / joblib / JSON / NetCDF and custom formats via export_model().

Installation

Core package (HPO only, minimal dependencies):

pip install het-ai

Full platform (HPO + DVC data versioning + MLflow tracking):

pip install "het-ai[platform]"

Install optional ML framework extras as needed:

# PyTorch + ONNX export
pip install "het-ai[torch]"

# TensorFlow / Keras / TFLite
pip install "het-ai[tensorflow]"

# scikit-learn
pip install "het-ai[sklearn]"

# Bayesian modelling with PyMC
pip install "het-ai[bayes]"

# Full platform + PyTorch
pip install "het-ai[platform,torch]"

# All dependencies required to run the bundled examples
pip install "het-ai[examples]"

Quick Start

Subclass BaseTrainer, declare your search space, and implement three hooks:

from het_ai.studio import BaseTrainer, TrainConfig
from het_ai.studio.types import TunableFloat, TunableInt


class MyTrainer(BaseTrainer):

    @BaseTrainer.search(
        lr=TunableFloat(1e-4, 1e-2, log=True),
        hidden=TunableInt(64, 512),
    )
    def train(self, data, lr, hidden):
        """Core training logic. Called once per Optuna trial."""
        model = build_model(hidden)
        score = fit(model, data, lr=lr)
        return score, model          # (metric value, exportable artifact)

    def load_data(self, dvc_data_root):
        """Load and return a DataBundle from the DVC data root."""
        ...

    def mock_data(self):
        """Return synthetic data for dry-run validation (no real data needed)."""
        ...

    def export_model(self, artifact, export_dir):
        """Persist the model artifact (e.g. ONNX, pickle) to export_dir."""
        ...


# ── Local validation (no production environment required) ────────────────
trainer = MyTrainer(TrainConfig(n_trials=2))
report  = trainer.dry_run()
print(report)
# {'score': 0.91, 'elapsed': 1.23, 'export_path': '/tmp/...'}

# ── Full HPO run (loads real data, launches Optuna study) ────────────────
result = trainer.run()

Data-Version-Driven Platform

When DVC and MLflow extras are installed, het-ai drives the entire pipeline from a data version tag to a registered model — fully automatically.

from het_ai.studio import BaseTrainer, TrainConfig
from het_ai.studio.bundle import DataBundle
from het_ai.dvc import DVCLoader, DVCConfig
from het_ai.mlflow import MLflowConfig
from het_ai.studio.types import TunableFloat
from pathlib import Path
import pandas as pd


class MyTrainer(BaseTrainer):

    @BaseTrainer.search(lr=TunableFloat(1e-4, 1e-2, log=True))
    def train(self, data: DataBundle, lr: float):
        ...

    def load_data(self, dvc_data_root: str) -> DataBundle:
        # Pull versioned data: GitHub tag → SeaweedFS → local
        loader = DVCLoader(self.config.dvc or DVCConfig())
        tag, sha = loader.pull(Path(dvc_data_root))

        df = pd.read_csv(f"{dvc_data_root}/data.csv")
        bundle = DataBundle(
            splits={"train": {"X": df[features].values,
                              "y": df["label"].values}},
            feature_list=features,
            target_list=["label"],
        )
        # Inject version metadata — MLflow will tag the Run automatically
        return loader.enrich_bundle(bundle, tag, sha)

    def mock_data(self) -> DataBundle: ...
    def export_model(self, artifact, export_dir) -> str: ...


config = TrainConfig(
    n_trials=100,
   dvc=DVCConfig(),        # reads DVC_GITHUB_* and SEAWEEDFS_* env vars
    mlflow=MLflowConfig(    # reads MLFLOW_TRACKING_URI and MLFLOW_EXPERIMENT
        experiment_name="my-project",
    ),
)

# DVC pull → Optuna HPO → model export →
# MLflow Run logged with dvc_version tag → model registered
result = MyTrainer(config).run()

After run() completes, every MLflow Run is tagged with the DVC version (e.g. dvc_version = release-20250525) that produced it, and the model is registered as my-project_prod in the MLflow Model Registry.

The minimum required environment variables for the full platform:

 # DVC / GitHub
 export DVC_GITHUB_REPO=org/data-repo
 export DVC_GITHUB_TOKEN=ghp_xxxxxxxxxxxx
export SEAWEEDFS_ENDPOINT=seaweedfs.internal:8333
export SEAWEEDFS_ACCESS_KEY=admin
export SEAWEEDFS_SECRET_KEY=change-me
export SEAWEEDFS_BUCKET=dvc-store

 # MLflow
 export MLFLOW_TRACKING_URI=http://mlflow.internal:5000
 export MLFLOW_EXPERIMENT=my-project

Multi-Objective Optimisation

Declare an objectives dict to enable Pareto-front HPO. The framework switches Optuna to multi-directional mode automatically:

from het_ai.studio.types import Result


class MultiObjectiveTrainer(BaseTrainer):

    objectives = {
        "accuracy":      "maximize",
        "model_size_mb": "minimize",
    }

    @BaseTrainer.search(lr=TunableFloat(1e-4, 1e-2))
    def train(self, data, lr):
        ...
        return Result(accuracy=acc, model_size_mb=size), model

Override select_best_trial(pareto_front) to apply a custom trade-off policy (default: highest value on the first objective).

TrainConfig Reference

All fields support environment-variable overrides. Precedence: env var > explicit argument > default value.

Supported Training Scenarios

The bundled test suite covers twelve representative scenarios, all runnable via dry_run() without any external dependencies:

Development & Testing

# Clone and install in editable mode with all development dependencies
git clone https://github.com/HeT-FTI/het-ai.git
cd het-ai
pip install -e ".[dev,examples]"

# Run the full test suite
pytest tests/ -v

# Run scenario smoke tests only
pytest tests/cases -v

The tests/cases directory currently contains 12 representative end-to-end dry-run scenarios for core workflow validation.

License

• Source code — Apache License 2.0

• Documentation — Creative Commons Attribution-NonCommercial 4.0 International (CC BY-NC 4.0)

For commercial use of the documentation, please contact HeT directly.

Copyright © 2026 Shenzhen HeT Intelligent Control Co., Ltd.