aitraining llm --train \
  --model google/gemma-3-270m \
  --data-path ./data.jsonl \
  --project-name sweep-experiment \
  --use-sweep \
  --sweep-backend optuna \
  --sweep-n-trials 20

from autotrain.trainers.clm.params import LLMTrainingParams

params = LLMTrainingParams(
    model="google/gemma-3-270m",
    data_path="./data.jsonl",
    project_name="sweep-experiment",

    # Enable sweep
    use_sweep=True,
    sweep_backend="optuna",
    sweep_n_trials=20,
    sweep_metric="eval_loss",
    sweep_direction="minimize",

    # Base parameters (sweep will vary some)
    trainer="sft",
    epochs=3,
    batch_size=4,
    lr=2e-5,
)

import json

# 字典格式（推荐）- 显式类型指定
sweep_params = json.dumps({
    "lr": {"type": "loguniform", "low": 1e-6, "high": 1e-3},
    "batch_size": {"type": "categorical", "values": [2, 4, 8]},
    "lora_r": {"type": "categorical", "values": [8, 16, 32, 64]},
    "warmup_ratio": {"type": "uniform", "low": 0.0, "high": 0.2},
    "epochs": {"type": "int", "low": 1, "high": 5},
})

# 列表格式（简写）- 仅用于分类值
sweep_params = json.dumps({
    "batch_size": [2, 4, 8],
    "lora_r": [8, 16, 32, 64],
})

params = LLMTrainingParams(
    ...
    use_sweep=True,
    sweep_params=sweep_params,  # JSON string
)

params = LLMTrainingParams(
    ...
    use_sweep=True,
    sweep_backend="optuna",
)

params = LLMTrainingParams(
    ...
    use_sweep=True,
    sweep_backend="grid",
)

params = LLMTrainingParams(
    ...
    use_sweep=True,
    sweep_backend="random",
)

params = LLMTrainingParams(
    ...
    use_sweep=True,
    sweep_metric="bleu",
    use_enhanced_eval=True,
    eval_metrics="bleu,rouge,bertscore",
    eval_batch_size=8,
)

import json

params = LLMTrainingParams(
    model="google/gemma-3-270m",
    data_path="./data.jsonl",
    project_name="lr-sweep",

    use_sweep=True,
    sweep_n_trials=10,
    sweep_params=json.dumps({
        "lr": {"type": "loguniform", "low": 1e-6, "high": 1e-3},
    }),

    # Fixed parameters
    trainer="sft",
    epochs=1,
    batch_size=4,
)

pip install optuna-dashboard
optuna-dashboard sqlite:///optuna.db

aitraining llm --train \
  --model google/gemma-3-270m \
  --data-path ./data.jsonl \
  --project-name sweep-experiment \
  --use-sweep \
  --sweep-backend optuna \
  --sweep-n-trials 20 \
  --log wandb \
  --wandb-sweep \
  --wandb-sweep-project my-sweep-project \
  --wandb-sweep-entity my-team

# 第一次运行创建 sweep（打印 "Created W&B sweep: abc123xyz"）
aitraining llm --train \
  --use-sweep --sweep-n-trials 10 \
  --wandb-sweep --wandb-sweep-project my-project

# 之后，继续同一个 sweep 添加更多 trial
aitraining llm --train \
  --use-sweep --sweep-n-trials 10 \
  --wandb-sweep --wandb-sweep-project my-project \
  --wandb-sweep-id abc123xyz

# 外部 W&B sweep 代理使用运行 ID 调用 AITraining
aitraining llm --train \
  --model google/gemma-3-270m \
  --data-path ./data.jsonl \
  --wandb-run-id $WANDB_RUN_ID \
  --lr $SWEEP_LR \
  --batch-size $SWEEP_BATCH_SIZE

aitraining llm --train \
  --model google/gemma-3-270m \
  --data-path ./data.jsonl \
  --project-name sweep-experiment \
  --use-sweep \
  --sweep-n-trials 10 \
  --post-trial-script 'echo "试验 $TRIAL_NUMBER 完成，指标 $TRIAL_METRIC_VALUE"'

aitraining llm --train \
  --use-sweep \
  --sweep-n-trials 20 \
  --post-trial-script 'if [ "$TRIAL_IS_BEST" = "true" ]; then git add . && git commit -m "最佳模型：试验 $TRIAL_NUMBER，指标 $TRIAL_METRIC_VALUE"; fi'

from autotrain.utils import HyperparameterSweep, SweepConfig, TrialInfo

def on_trial_complete(trial_info: TrialInfo):
    """每次试验完成后调用。"""
    print(f"试验 {trial_info.trial_number} 完成")
    print(f"  参数: {trial_info.params}")
    print(f"  指标: {trial_info.metric_value}")
    print(f"  是最佳: {trial_info.is_best}")

    if trial_info.is_best:
        # 为最佳试验执行特殊操作
        save_best_model(trial_info.output_dir)

config = SweepConfig(
    parameters={"lr": (1e-5, 1e-3, "log_uniform")},
    n_trials=10,
    backend="optuna",
    post_trial_callback=on_trial_complete,
)

sweep = HyperparameterSweep(config, train_function)
result = sweep.run()