scripts/train_unsloth.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "datasets>=2.21.0",
#   "matplotlib>=3.9.0",
#   "pandas>=2.2.2",
#   "peft>=0.14.0",
#   "trl>=0.13.0",
#   "transformers>=4.46.0",
#   "unsloth>=2025.1.0",
# ]
# ///

"""BDO.ai GRPO fine-tuning pipeline."""

from __future__ import annotations

import argparse
import json
import os
import random
import sys
from pathlib import Path
from statistics import mean
from typing import Any, Callable

try:
    import unsloth  # noqa: F401
    from unsloth import FastLanguageModel
    import torch
    import matplotlib.pyplot as plt
    from datasets import Dataset
    from trl import GRPOConfig, GRPOTrainer
except ImportError:
    print("Warning: unsloth/trl/transformers/matplotlib are not installed.")
    print("Run this script on a GPU machine or with `hf jobs uv run`.")
    sys.exit(1)

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from bdo_ai_env.training import build_prompt
from models import BDOAction
from server.bdo_environment import BDOEnvironment


def balanced_policy(observation: dict[str, Any]) -> dict[str, Any]:
    highest = max(observation["nodes"], key=lambda node: node["reported_demand"])
    weakest_signal = min(observation["nodes"], key=lambda node: node["biometric_signal"])
    queue = observation["high_risk_queue"]

    actions: list[dict[str, Any]] = []
    if queue and weakest_signal["biometric_signal"] < 0.7:
        actions.append(
            {"name": "trigger_field_audit", "params": {"village": weakest_signal["village"]}}
        )
    elif weakest_signal["biometric_signal"] < 0.55:
        actions.append(
            {"name": "dispatch_repair", "params": {"village": weakest_signal["village"]}}
        )

    spend = min(
        observation["treasury"]["district_budget"],
        max(2000, int(highest["reported_demand"] * 0.85)),
    )
    actions.append(
        {"name": "allocate_funds", "params": {"village": highest["village"], "amount": spend}}
    )
    actions.append(
        {"name": "approve_batch", "params": {"village": highest["village"], "mode": "conservative"}}
    )
    if queue:
        actions.append({"name": "reject_transfer", "params": {"transfer_id": queue[0]["transfer_id"]}})

    avg_signal = sum(node["biometric_signal"] for node in observation["nodes"]) / len(observation["nodes"])
    predicted_fraud = round(min(0.9, max(0.12, 1.0 - avg_signal)), 3)
    return {
        "predicted_fraud_level": predicted_fraud,
        "thought_process": (
            f"{highest['village']} has the highest reported demand while "
            f"{weakest_signal['village']} has the weakest biometrics. "
            "Use conservative approvals and handle the riskiest transfer."
        ),
        "actions": actions,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run GRPO fine-tuning for BDO.ai.")
    parser.add_argument("--model-name", default="Qwen/Qwen2.5-1.5B-Instruct")
    parser.add_argument("--max-seq-length", type=int, default=4096)
    parser.add_argument("--dataset-episodes", type=int, default=24)
    parser.add_argument("--eval-episodes", type=int, default=20)
    parser.add_argument("--steps", type=int, default=120)
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=1e-5)
    parser.add_argument("--num-generations", type=int, default=4)
    parser.add_argument("--max-prompt-length", type=int, default=3072)
    parser.add_argument("--max-completion-length", type=int, default=640)
    parser.add_argument("--output-dir", default="artifacts/grpo_outputs")
    parser.add_argument("--seed", type=int, default=3407)
    return parser.parse_args()


def collect_grpo_prompts(*, dataset_episodes: int, seed: int) -> list[dict[str, str]]:
    scenarios = ["calm_year", "black_swan", "fraud_syndicate", "rapid_migration"]
    rows: list[dict[str, str]] = []

    for episode in range(dataset_episodes):
        scenario = scenarios[episode % len(scenarios)]
        env = BDOEnvironment(scenario=scenario, seed=seed + episode)
        observation = env.reset(seed=seed + episode, scenario=scenario)
        done = observation.done

        while not done:
            prompt = build_prompt(observation.model_dump(mode="json", exclude_none=True))
            rows.append({"prompt": prompt, "scenario": scenario})
            observation = env.step(BDOAction.model_validate(balanced_policy(observation.model_dump(mode="json", exclude_none=True))))
            done = observation.done

    Path("artifacts").mkdir(exist_ok=True)
    Path("artifacts/grpo_prompt_preview.json").write_text(json.dumps(rows[:40], indent=2), encoding="utf-8")
    return rows


def extract_json_object(text: str) -> str:
    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ValueError("Model output did not contain a JSON object.")
    return text[start : end + 1]


def clip_prompt(
    prompt: str,
    tokenizer: Any,
    *,
    max_prompt_tokens: int,
) -> str:
    encoded = tokenizer(
        prompt,
        add_special_tokens=False,
        truncation=True,
        max_length=max_prompt_tokens,
    )
    return tokenizer.decode(encoded["input_ids"], skip_special_tokens=False)


def completion_quality_penalty(text: str) -> float:
    penalty = -2.0
    if "{" in text:
        penalty += 0.5
    if "}" in text:
        penalty += 0.35
    if '"actions"' in text:
        penalty += 0.35
    if '"predicted_fraud_level"' in text:
        penalty += 0.3
    if '"thought_process"' in text:
        penalty += 0.25
    return min(-0.2, penalty)


def configure_generation_defaults(model: Any, tokenizer: Any) -> None:
    model.generation_config.max_length = None
    model.generation_config.max_new_tokens = None
    model.generation_config.pad_token_id = tokenizer.eos_token_id


def normalize_completion(completion: Any) -> str:
    if isinstance(completion, str):
        return completion
    if isinstance(completion, list):
        chunks: list[str] = []
        for item in completion:
            if isinstance(item, dict):
                chunks.append(str(item.get("content", "")))
            else:
                chunks.append(str(item))
        return "".join(chunks)
    return str(completion)


def build_reward_function(seed: int) -> Callable[..., list[float]]:
    scenarios = ["black_swan", "fraud_syndicate", "rapid_migration"]

    def reward_func(prompts, completions, **kwargs):
        rewards: list[float] = []
        for idx, completion in enumerate(completions):
            scenario = scenarios[idx % len(scenarios)]
            env = BDOEnvironment(scenario=scenario, seed=seed + idx)
            env.reset(seed=seed + idx, scenario=scenario)
            completion_text = normalize_completion(completion)
            try:
                action = BDOAction.model_validate_json(extract_json_object(completion_text))
                observation = env.step(action)
                info = observation.info or observation.metadata
                reward = float(info["training_reward"])
            except Exception:
                reward = completion_quality_penalty(completion_text)
            rewards.append(reward)
        return rewards

    return reward_func


def model_policy(
    model: Any,
    tokenizer: Any,
    *,
    max_seq_length: int,
    max_new_tokens: int = 512,
) -> Callable[[dict[str, Any]], dict[str, Any]]:
    if hasattr(FastLanguageModel, "for_inference"):
        FastLanguageModel.for_inference(model)

    def _policy(observation: dict[str, Any]) -> dict[str, Any]:
        prompt = clip_prompt(
            build_prompt(observation),
            tokenizer,
            max_prompt_tokens=max(256, max_seq_length - max_new_tokens),
        )
        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=max_seq_length - max_new_tokens,
        ).to(model.device)
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
        )
        completion = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)
        try:
            return BDOAction.model_validate_json(extract_json_object(completion)).model_dump(mode="json", exclude_none=True)
        except Exception:
            return balanced_policy(observation)

    return _policy


def evaluate_policy(
    policy_fn: Callable[[dict[str, Any]], dict[str, Any]],
    *,
    episodes: int,
    seed: int,
) -> list[dict[str, Any]]:
    scenarios = ["black_swan", "fraud_syndicate", "rapid_migration"]
    rows: list[dict[str, Any]] = []
    for episode in range(episodes):
        scenario = scenarios[episode % len(scenarios)]
        env = BDOEnvironment(scenario=scenario, seed=seed + episode)
        observation = env.reset(seed=seed + episode, scenario=scenario)
        total_reward = 0.0
        total_training_reward = 0.0
        belief_scores: list[float] = []
        while not observation.done:
            action_payload = policy_fn(observation.model_dump(mode="json", exclude_none=True))
            observation = env.step(BDOAction.model_validate(action_payload))
            info = observation.info or observation.metadata
            total_reward += float(observation.reward or 0.0)
            total_training_reward += float(info.get("training_reward", 0.0))
            belief_scores.append(float(info["reward_breakdown"]["belief_accuracy"]))
        rows.append(
            {
                "episode": episode,
                "scenario": scenario,
                "total_reward": round(total_reward, 4),
                "total_training_reward": round(total_training_reward, 4),
                "avg_belief_accuracy": round(mean(belief_scores) if belief_scores else 0.0, 4),
            }
        )
    return rows


def write_reward_artifacts(results: dict[str, list[dict[str, Any]]]) -> dict[str, Any]:
    artifacts_dir = Path("artifacts")
    artifacts_dir.mkdir(exist_ok=True)
    summary = {}
    for label, rows in results.items():
        scenario_summary: dict[str, dict[str, float | int]] = {}
        scenarios = sorted({row["scenario"] for row in rows})
        for scenario in scenarios:
            scenario_rows = [row for row in rows if row["scenario"] == scenario]
            scenario_summary[scenario] = {
                "episodes": len(scenario_rows),
                "mean_total_reward": round(mean(row["total_reward"] for row in scenario_rows), 4),
                "mean_total_training_reward": round(
                    mean(row["total_training_reward"] for row in scenario_rows), 4
                ),
                "mean_belief_accuracy": round(
                    mean(row["avg_belief_accuracy"] for row in scenario_rows), 4
                ),
            }
        summary[label] = {
            "episodes": rows,
            "mean_total_reward": round(mean(row["total_reward"] for row in rows), 4),
            "mean_total_training_reward": round(mean(row["total_training_reward"] for row in rows), 4),
            "mean_belief_accuracy": round(mean(row["avg_belief_accuracy"] for row in rows), 4),
            "scenario_summary": scenario_summary,
        }

    (artifacts_dir / "grpo_reward_curve.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    plt.figure(figsize=(10, 5))
    for label, rows in results.items():
        plt.plot([row["episode"] for row in rows], [row["total_reward"] for row in rows], label=label)
    plt.xlabel("evaluation episode")
    plt.ylabel("total environment reward")
    plt.title("BDO.ai GRPO Reward Comparison")
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(artifacts_dir / "grpo_reward_curve.png", dpi=180)
    plt.close()

    labels = list(summary.keys())
    mean_rewards = [summary[label]["mean_total_reward"] for label in labels]
    plt.figure(figsize=(7, 5))
    plt.bar(labels, mean_rewards, color=["#4C78A8", "#F58518"][: len(labels)])
    plt.ylabel("mean total environment reward")
    plt.title("BDO.ai GRPO Mean Reward")
    plt.grid(True, axis="y", alpha=0.3)
    plt.tight_layout()
    plt.savefig(artifacts_dir / "grpo_mean_reward.png", dpi=180)
    plt.close()
    return summary


def write_train_metrics(log_history: list[dict[str, Any]]) -> None:
    rows = []
    for row in log_history:
        if "step" in row and any(key in row for key in ("reward", "loss", "kl")):
            clean = {"step": int(row["step"])}
            for key in ("reward", "loss", "kl", "entropy"):
                if key in row:
                    clean[key] = float(row[key])
            rows.append(clean)
    Path("artifacts").mkdir(exist_ok=True)
    Path("artifacts/grpo_training_metrics.json").write_text(json.dumps(rows, indent=2), encoding="utf-8")
    if rows and any("reward" in row for row in rows):
        reward_rows = [row for row in rows if "reward" in row]
        plt.figure(figsize=(8, 5))
        plt.plot([row["step"] for row in reward_rows], [row["reward"] for row in reward_rows])
        plt.xlabel("training step")
        plt.ylabel("grpo reward")
        plt.title("BDO.ai GRPO Reward Curve")
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig("artifacts/grpo_training_curve.png", dpi=180)
        plt.close()


def main() -> None:
    args = parse_args()
    random.seed(args.seed)

    print("Building environment-linked GRPO prompts...")
    dataset_rows = collect_grpo_prompts(dataset_episodes=args.dataset_episodes, seed=args.seed)
    train_dataset = Dataset.from_list(dataset_rows)
    print(f"Prepared {len(train_dataset)} GRPO prompts.")

    print(f"Loading model: {args.model_name}")
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=args.model_name,
        max_seq_length=args.max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )
    configure_generation_defaults(model, tokenizer)

    print("Running pre-training evaluation...")
    untrained_policy = model_policy(
        model,
        tokenizer,
        max_seq_length=args.max_prompt_length + args.max_completion_length,
        max_new_tokens=args.max_completion_length,
    )
    untrained_results = evaluate_policy(untrained_policy, episodes=args.eval_episodes, seed=args.seed + 100)

    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        lora_alpha=16,
        use_gradient_checkpointing="unsloth",
        random_state=args.seed,
    )

    trainer = GRPOTrainer(
        model=model,
        processing_class=tokenizer,
        reward_funcs=build_reward_function(args.seed + 900),
        args=GRPOConfig(
            output_dir=args.output_dir,
            learning_rate=args.learning_rate,
            per_device_train_batch_size=args.batch_size,
            gradient_accumulation_steps=args.grad_accum,
            num_generations=args.num_generations,
            max_prompt_length=args.max_prompt_length,
            max_completion_length=args.max_completion_length,
            max_steps=args.steps,
            logging_steps=1,
            save_strategy="no",
            report_to=[],
            bf16=bool(torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
            fp16=not bool(torch.cuda.is_available() and torch.cuda.is_bf16_supported()),
        ),
        train_dataset=train_dataset,
    )

    print("Starting GRPO training...")
    train_result = trainer.train()

    adapter_dir = Path("artifacts/grpo_model")
    adapter_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)
    write_train_metrics(trainer.state.log_history)

    # Upload adapter to HF Space so the UI can load it on restart
    hf_token = os.environ.get("HF_TOKEN")
    space_id = os.environ.get("SPACE_ID")  # set automatically inside HF jobs
    if hf_token and space_id:
        try:
            from huggingface_hub import HfApi
            print(f"Uploading adapter to Space {space_id} ...")
            HfApi(token=hf_token).upload_folder(
                folder_path=str(adapter_dir),
                repo_id=space_id,
                repo_type="space",
                path_in_repo="artifacts/grpo_model",
            )
            print("Adapter uploaded. UI will use LLM policy after Space restart.")
        except Exception as _upload_err:
            print(f"Warning: adapter upload failed ({_upload_err}). Weights saved locally only.")

    print("Running post-training evaluation...")
    trained_policy = model_policy(
        model,
        tokenizer,
        max_seq_length=args.max_prompt_length + args.max_completion_length,
        max_new_tokens=args.max_completion_length,
    )
    reward_results = {
        "untrained_qwen_base": untrained_results,
        "trained_qwen_grpo": evaluate_policy(trained_policy, episodes=args.eval_episodes, seed=args.seed + 300),
    }
    reward_summary = write_reward_artifacts(reward_results)

    summary = {
        "model_name": args.model_name,
        "dataset_prompts": len(train_dataset),
        "dataset_episodes": args.dataset_episodes,
        "eval_episodes": args.eval_episodes,
        "grpo_steps": args.steps,
        "train_runtime_seconds": round(float(train_result.metrics.get("train_runtime", 0.0)), 2),
        "reward_summary": reward_summary,
        "artifacts": {
            "grpo_training_metrics_json": "artifacts/grpo_training_metrics.json",
            "grpo_training_curve_png": "artifacts/grpo_training_curve.png",
            "grpo_reward_curve_json": "artifacts/grpo_reward_curve.json",
            "grpo_reward_curve_png": "artifacts/grpo_reward_curve.png",
            "grpo_mean_reward_png": "artifacts/grpo_mean_reward.png",
            "grpo_summary_json": "artifacts/grpo_training_summary.json",
            "adapter_dir": str(adapter_dir),
        },
    }
    Path("artifacts/grpo_training_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()