chore: import upstream snapshot with attribution

examples/spider/train_sql_agent.py

@@ -0,0 +1,217 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Train an SQL agent on the Spider dataset using Agent-lightning.
+
+This module provides a training script for SQL agents using different model configurations.
+The script supports three different training configurations:
+
+1. 'fast' - A lightweight configuration optimized for CI testing with reduced epochs
+2. 'qwen' - Standard configuration using Qwen-2.5-Coder-1.5B-Instruct model
+3. 'llama' - Configuration using LLaMA-3.2-1B-Instruct model with JSON formatting
+
+Usage:
+    python train_sql_agent.py fast    # Fast training for CI/testing
+    python train_sql_agent.py qwen    # Standard Qwen model training
+    python train_sql_agent.py llama   # LLaMA model training
+
+The script uses reinforcement learning with VERL framework
+to train agents on the Spider dataset for text-to-SQL generation tasks.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, Optional
+
+import pandas as pd
+from sql_agent import LitSQLAgent
+
+import agentlightning as agl
+
+RL_TRAINING_CONFIG: Dict[str, Any] = {
+    "algorithm": {
+        "adv_estimator": "grpo",
+        "use_kl_in_reward": False,
+    },
+    "data": {
+        "train_files": "data/train_spider.parquet",
+        "val_files": "data/test_dev_500.parquet",
+        "train_batch_size": 32,
+        "max_prompt_length": 4096,
+        "max_response_length": 2048,
+        "truncation": "error",
+    },
+    "actor_rollout_ref": {
+        "rollout": {
+            "tensor_model_parallel_size": 1,
+            "n": 4,
+            "log_prob_micro_batch_size_per_gpu": 4,
+            "multi_turn": {"format": "hermes"},
+            "name": "vllm",
+            "gpu_memory_utilization": 0.8,
+            "engine_kwargs": {
+                "vllm": {
+                    "enable_auto_tool_choice": True,
+                    "tool_call_parser": "hermes",
+                }
+            },
+        },
+        "actor": {
+            "ppo_mini_batch_size": 32,
+            "ppo_micro_batch_size_per_gpu": 4,
+            "optim": {"lr": 1e-6},
+            "use_kl_loss": False,
+            "kl_loss_coef": 0.0,
+            "entropy_coeff": 0,
+            "clip_ratio_low": 0.2,
+            "clip_ratio_high": 0.3,
+            "fsdp_config": {
+                "param_offload": True,
+                "optimizer_offload": True,
+            },
+        },
+        "ref": {
+            "log_prob_micro_batch_size_per_gpu": 8,
+            "fsdp_config": {"param_offload": True},
+        },
+        "model": {
+            "path": "Qwen/Qwen2.5-Coder-1.5B-Instruct",
+            "use_remove_padding": True,
+            "enable_gradient_checkpointing": True,
+        },
+    },
+    "trainer": {
+        "n_gpus_per_node": 1,
+        "val_before_train": True,
+        "critic_warmup": 0,
+        "logger": ["console", "wandb"],
+        "project_name": "AgentLightning",
+        "experiment_name": "spider",
+        "nnodes": 1,
+        "test_freq": 32,
+        "total_epochs": 2,
+    },
+}
+
+
+def config_train_fast() -> Dict[str, Any]:
+    """A fast training run for CI testing purposes."""
+
+    # `EXPERIMENT_NAME="spider_$(date +%Y%m%d%H%M%S)"`
+    timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+    EXPERIMENT_NAME = f"spider_{timestamp}"
+
+    # `PROJECT_NAME=AgentLightningCI`
+    PROJECT_NAME = "AgentLightningCI"
+
+    # Simulate writing to $GITHUB_OUTPUT if it’s set
+    github_output = os.getenv("GITHUB_OUTPUT")
+    if github_output:
+        with open(github_output, "a") as f:
+            f.write(f"project_name={PROJECT_NAME}\n")
+            f.write(f"run_name={EXPERIMENT_NAME}\n")
+
+    print("Set environment variables:")
+    print(f"PROJECT_NAME={PROJECT_NAME}")
+    print(f"EXPERIMENT_NAME={EXPERIMENT_NAME}")
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["gpu_memory_utilization"] = 0.6
+    config["actor_rollout_ref"]["model"]["path"] = "Qwen/Qwen2.5-Coder-0.5B-Instruct"
+    config["data"]["val_files"] = "data/test_dev.parquet"
+    config["trainer"]["total_epochs"] = 1
+    config["trainer"]["total_training_steps"] = 1
+    config["trainer"]["experiment_name"] = EXPERIMENT_NAME
+    config["trainer"]["project_name"] = PROJECT_NAME
+    config["trainer"]["test_freq"] = 1
+    return config
+
+
+def config_train_qwen() -> Dict[str, Any]:
+    """A configuration for training with Qwen-2.5B."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    return config
+
+
+def config_train_npu() -> Dict[str, Any]:
+    """A configuration for training with NPU."""
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    del config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["enable_auto_tool_choice"]
+    del config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["tool_call_parser"]
+    del config["trainer"]["logger"][1]
+    config["actor_rollout_ref"]["actor"]["use_torch_compile"] = False
+    config["trainer"]["val_before_train"] = False
+    config["trainer"]["save_freq"] = 256
+    config["trainer"]["device"] = "npu"
+    return config
+
+
+def config_train_llama() -> Dict[str, Any]:
+    """A configuration for training with LLaMA-3.2-1B-Instruct.
+
+    You will need a `HF_TOKEN` set to run with this config.
+    """
+
+    config = deepcopy(RL_TRAINING_CONFIG)
+    config["actor_rollout_ref"]["rollout"]["multi_turn"]["format"] = "llama3_json"
+    config["actor_rollout_ref"]["rollout"]["engine_kwargs"]["vllm"]["tool_call_parser"] = "llama3_json"
+    config["actor_rollout_ref"]["model"]["path"] = "meta-llama/Llama-3.2-1B-Instruct"
+    return config
+
+
+def train(config: Dict[str, Any], active_agent: Optional[str]) -> None:
+    """Train the SQL agent with the given configuration."""
+
+    agent = LitSQLAgent()
+    algorithm = agl.VERL(config)
+    trainer = agl.Trainer(n_runners=10, algorithm=algorithm, adapter={"agent_match": active_agent})
+    print("Adapter agent match acknowledged:", trainer.adapter.agent_match)  # type: ignore
+
+    train_data = pd.read_parquet(config["data"]["train_files"]).to_dict(orient="records")  # type: ignore
+    val_data = pd.read_parquet(config["data"]["val_files"]).to_dict(orient="records")  # type: ignore
+    trainer.fit(agent, train_dataset=train_data, val_dataset=val_data)  # type: ignore
+
+
+def main() -> None:
+    """Main function to parse arguments and run training."""
+    parser = argparse.ArgumentParser(
+        description="Train an SQL agent on the Spider dataset using different model configurations"
+    )
+
+    parser.add_argument(
+        "config",
+        choices=["fast", "qwen", "llama", "npu"],
+        help="Training configuration: 'fast' (CI testing), 'qwen' (Qwen-2.5-Coder-1.5B), 'llama' (LLaMA-3.2-3B),'npu' (Train with NPU)",
+    )
+
+    parser.add_argument(
+        "--active-agent", type=str, help="Override the active agent name (default: auto-generated based on config)"
+    )
+
+    args = parser.parse_args()
+
+    # Get the appropriate configuration
+    config_functions = {
+        "fast": config_train_fast,
+        "qwen": config_train_qwen,
+        "llama": config_train_llama,
+        "npu": config_train_npu,
+    }
+    config = config_functions[args.config]()
+
+    # Set active agent - use provided value or default based on config choice
+    active_agent = args.active_agent
+
+    print(f"Starting training with '{args.config}' configuration...")
+    print(f"Active agent: {active_agent}")
+
+    train(config, active_agent)
+
+
+if __name__ == "__main__":
+    main()