chore: import upstream snapshot with attribution

examples/multi_task/train_multi_task.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python
+"""Multi-Task Learning with Loss Balancing in Ludwig.
+
+Trains four models on the UCI Wine Quality dataset with two output features:
+  - quality_score  : raw 0-10 quality score (regression)
+  - quality_binary : quality >= 7 is "good" (binary classification)
+
+Compares loss balancing strategies:
+  1. none         — static weighted sum (baseline)
+  2. famo         — Fast Adaptive Multitask Optimization (available now)
+  3. uncertainty  — Homoscedastic uncertainty weighting (available now)
+  4. nash_mtl     — Nash bargaining solution (requires PR #4092)
+
+# Colab: !pip install ludwig
+
+Usage:
+    python train_multi_task.py
+"""
+
+import logging
+import os
+import shutil
+import warnings
+
+import pandas as pd
+
+logging.basicConfig(level=logging.WARNING)
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+WINE_URL = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
+
+WINE_FEATURES = [
+    "fixed_acidity",
+    "volatile_acidity",
+    "citric_acid",
+    "residual_sugar",
+    "chlorides",
+    "free_sulfur_dioxide",
+    "total_sulfur_dioxide",
+    "density",
+    "pH",
+    "sulphates",
+    "alcohol",
+]
+
+
+def load_dataset() -> pd.DataFrame:
+    """Download and prepare the dual-output wine quality dataset."""
+    print("Downloading wine quality dataset...")
+    df = pd.read_csv(WINE_URL, sep=";")
+    df.columns = [c.replace(" ", "_") for c in df.columns]
+    # quality_score: keep the raw 0-10 numerical score
+    df["quality_score"] = df["quality"].astype(float)
+    # quality_binary: 1 if quality >= 7 (good wine), else 0
+    df["quality_binary"] = (df["quality"] >= 7).astype(int)
+    df = df.drop(columns=["quality"])
+    print(f"  {len(df)} rows | good wines (quality >= 7): {df['quality_binary'].mean():.1%}")
+    print(f"  quality_score range: {df['quality_score'].min():.0f} – {df['quality_score'].max():.0f}")
+    return df
+
+
+# ---------------------------------------------------------------------------
+# Ludwig config helpers
+# ---------------------------------------------------------------------------
+
+
+def _input_features() -> list:
+    return [{"name": feat, "type": "number", "preprocessing": {"normalization": "zscore"}} for feat in WINE_FEATURES]
+
+
+def _base_config(loss_balancing: str) -> dict:
+    return {
+        "model_type": "ecd",
+        "input_features": _input_features(),
+        "output_features": [
+            {"name": "quality_score", "type": "number"},
+            {"name": "quality_binary", "type": "binary"},
+        ],
+        "combiner": {
+            "type": "concat",
+            "num_fc_layers": 2,
+            "output_size": 128,
+            "dropout": 0.1,
+        },
+        "trainer": {
+            "epochs": 30,
+            "learning_rate": 0.001,
+            "batch_size": 128,
+            "loss_balancing": loss_balancing,
+        },
+    }
+
+
+# ---------------------------------------------------------------------------
+# Training helper
+# ---------------------------------------------------------------------------
+
+
+def train_and_evaluate(
+    name: str,
+    config: dict,
+    dataset: pd.DataFrame,
+    output_dir: str,
+) -> dict | None:
+    """Train a Ludwig model and return evaluation metrics.
+
+    Returns a dict with metric values, or None if training failed.
+    """
+    from ludwig.api import LudwigModel
+
+    result_dir = os.path.join(output_dir, name)
+    shutil.rmtree(result_dir, ignore_errors=True)
+
+    print(f"\n--- Training: {name} ---")
+    try:
+        model = LudwigModel(config=config, logging_level=logging.WARNING)
+        result = model.train(
+            dataset=dataset,
+            experiment_name="multi_task",
+            model_name=name,
+            output_directory=result_dir,
+        )
+
+        # Extract final validation metrics
+        metrics = {}
+        vset = result.train_stats.validation or {}
+        # quality_score: mean absolute error (lower is better)
+        score_metrics = vset.get("quality_score", {})
+        metrics["score_mae"] = _last_value(score_metrics.get("mean_absolute_error", []))
+        # quality_binary: ROC AUC (higher is better)
+        binary_metrics = vset.get("quality_binary", {})
+        metrics["binary_roc_auc"] = _last_value(binary_metrics.get("roc_auc", []))
+        return metrics
+
+    except Exception as exc:
+        warnings.warn(f"Training '{name}' failed: {exc}", stacklevel=2)
+        return None
+
+
+def _last_value(series) -> float | None:
+    """Return the last numeric value in a list, or None."""
+    if not series:
+        return None
+    val = series[-1]
+    if isinstance(val, (list, tuple)):
+        val = val[-1]
+    try:
+        return float(val)
+    except (TypeError, ValueError):
+        return None
+
+
+# ---------------------------------------------------------------------------
+# Comparison table
+# ---------------------------------------------------------------------------
+
+
+def print_comparison_table(results: dict) -> None:
+    """Print a formatted side-by-side comparison of all methods."""
+    col_w = 14
+    header = f"{'Method':<{col_w}} | {'Score MAE':>{col_w}} | {'Binary ROC-AUC':>{col_w}}"
+    separator = "-" * len(header)
+    print()
+    print("=" * len(header))
+    print("  Multi-Task Loss Balancing — Comparison")
+    print("=" * len(header))
+    print(header)
+    print(separator)
+    for method, metrics in results.items():
+        if metrics is None:
+            mae_str = "FAILED"
+            auc_str = "FAILED"
+        else:
+            mae = metrics.get("score_mae")
+            auc = metrics.get("binary_roc_auc")
+            mae_str = f"{mae:.4f}" if mae is not None else "n/a"
+            auc_str = f"{auc:.4f}" if auc is not None else "n/a"
+        print(f"{method:<{col_w}} | {mae_str:>{col_w}} | {auc_str:>{col_w}}")
+    print(separator)
+    print("  Score MAE: lower is better | Binary ROC-AUC: higher is better")
+    print()
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main():
+    output_dir = "./results"
+    os.makedirs(output_dir, exist_ok=True)
+
+    df = load_dataset()
+
+    # Methods to compare. nash_mtl is attempted but skipped gracefully if
+    # PR #4092 is not yet merged.
+    methods = [
+        ("none", False),
+        ("famo", False),
+        ("uncertainty", False),
+        ("nash_mtl", True),  # requires PR #4092
+    ]
+
+    results = {}
+    for method, requires_pr in methods:
+        if requires_pr:
+            print(f"\n--- Skipping {method} (requires PR #4092 / Ludwig >= 0.14) ---")
+            print("  To enable, install Ludwig from the 'future-capabilities' branch:")
+            print("    pip install git+https://github.com/ludwig-ai/ludwig@future-capabilities")
+            results[method] = None
+            continue
+
+        config = _base_config(method)
+        results[method] = train_and_evaluate(method, config, df, output_dir)
+
+    # Attempt nash_mtl — will succeed if PR #4092 is available
+    try:
+        from ludwig.api import LudwigModel
+
+        config = _base_config("nash_mtl")
+        # Try instantiating to check if nash_mtl is a valid option
+        model = LudwigModel(config=config, logging_level=logging.WARNING)
+        del model
+        print("\n  nash_mtl is available — training now...")
+        results["nash_mtl"] = train_and_evaluate("nash_mtl", config, df, output_dir)
+    except Exception:
+        pass  # already marked as None above
+
+    print_comparison_table(results)
+
+
+if __name__ == "__main__":
+    main()