scripts/hf_eval_supplymind_adapters.py

# /// script
# dependencies = [
#   "torch",
#   "transformers>=4.45.0",
#   "peft>=0.13.0",
#   "accelerate",
#   "huggingface_hub",
#   "pydantic",
#   "pyyaml",
# ]
# ///
from __future__ import annotations

import argparse
from collections import Counter
import json
import os
import re
import sys
import time
from pathlib import Path
from statistics import mean
from typing import Any

import torch
from huggingface_hub import snapshot_download
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig


REPO_ID = "rishavutk/supplymind"
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"


def log(message: str, **fields: Any) -> None:
    print(json.dumps({"message": message, **fields}, sort_keys=True), flush=True)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--role", choices=["center", "warehouse"], required=True)
    parser.add_argument("--adapter-id", default="")
    parser.add_argument("--sft-adapter-id", default="")
    parser.add_argument("--grpo-adapter-id", default="")
    parser.add_argument("--task-id", default="v2_train_easy")
    parser.add_argument("--seeds", default="101,113,127")
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--model-id", default=MODEL_ID)
    parser.add_argument("--load-in-4bit", action="store_true")
    return parser.parse_args()


def prepare_repo() -> None:
    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
    local_dir = Path(
        snapshot_download(
            repo_id=REPO_ID,
            repo_type="space",
            allow_patterns=["src/**", "configs/**"],
            local_dir=Path("supplymind_snapshot").resolve(),
        )
    )
    os.environ["SUPPLYMIND_REWARD_CONFIG"] = str(local_dir / "configs" / "supplymind_v2_rewards.yaml")
    sys.path.insert(0, str(local_dir / "src"))
    log("repo_ready", path=str(local_dir), config_exists=Path(os.environ["SUPPLYMIND_REWARD_CONFIG"]).exists())


def extract_json(text: str) -> dict[str, Any] | None:
    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(0))
    except json.JSONDecodeError:
        return None


def compact_observation(observation: Any, role: str, warehouse_id: str | None = None) -> dict[str, Any]:
    data = observation.model_dump(mode="json")
    data["scenario_info"].pop("public_rules", None)
    if role == "center":
        return {
            "role": "center",
            "round_index": data["round_index"],
            "remaining_rounds": data["remaining_rounds"],
            "task_id": data["task_id"],
            "scenario_info": data["scenario_info"],
            "center": data["center"],
            "warehouse_summaries": data["center"].get("warehouse_summaries", []),
            "market_signals": data["center"].get("market_signals", []),
            "feedback": data.get("feedback", {}),
        }
    if warehouse_id:
        warehouse = data["warehouses"][warehouse_id]
        return {
            "role": "warehouse",
            "warehouse_id": warehouse_id,
            "round_index": data["round_index"],
            "remaining_rounds": data["remaining_rounds"],
            "task_id": data["task_id"],
            "scenario_info": data["scenario_info"],
            "warehouse": warehouse,
            "pending_transfer_proposals": warehouse.get("pending_transfer_proposals", []),
            "feedback": data.get("feedback", {}),
        }
    return {
        "role": "warehouse",
        "round_index": data["round_index"],
        "remaining_rounds": data["remaining_rounds"],
        "task_id": data["task_id"],
        "scenario_info": data["scenario_info"],
        "warehouses": data["warehouses"],
        "pending_transfer_proposals": data["center"].get("pending_transfer_proposals", []),
        "feedback": data.get("feedback", {}),
    }


def system_prompt(role: str) -> str:
    if role == "center":
        return (
            "You are the center policy in SupplyMind. Return only strict JSON matching CenterAction: "
            "central_procurements, central_liquidations, central_replenishments, inventory_transfer_proposals, offer_matches. "
            "Warehouses are controlled by a fixed heuristic. Earn margin and a small share of realized service profit, "
            "but avoid waste, stockouts, overpriced actions, and needless shipments. Empty lists are only appropriate "
            "when no useful procurement, liquidation, replenishment, transfer proposal, or offer match exists."
        )
    return (
        "You are the shared warehouse policy in SupplyMind, copied across all warehouses. "
        "You control exactly one warehouse from the user observation. Return only strict JSON matching WarehouseAction: "
        "order_decisions, inventory_offers, inventory_requests, transfer_responses, and local_priority. "
        "The center is controlled by a fixed heuristic. Accept orders you can serve, request needed stock, "
        "and reject bad or impossible commitments. Only use order_id and proposal_id values visible in this observation. "
        "Do not invent IDs, do not use markdown, and prefer fewer high-confidence actions over broad noisy actions."
    )


def generate_action(
    model: Any,
    tokenizer: Any,
    role: str,
    observation: Any,
    max_new_tokens: int,
    warehouse_id: str | None = None,
) -> tuple[dict[str, Any] | None, str]:
    prompt = [
        {"role": "system", "content": system_prompt(role)},
        {"role": "user", "content": json.dumps(compact_observation(observation, role, warehouse_id), separators=(",", ":"))},
    ]
    text = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(text, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id,
        )
    generated = tokenizer.decode(output[0][inputs["input_ids"].shape[-1] :], skip_special_tokens=True)
    return extract_json(generated), generated


def _len_list(payload: dict[str, Any], key: str) -> int:
    value = payload.get(key, [])
    return len(value) if isinstance(value, list) else 0


def center_action_stats(payload: dict[str, Any] | None) -> dict[str, int]:
    payload = payload or {}
    return {
        "central_procurements": _len_list(payload, "central_procurements"),
        "central_liquidations": _len_list(payload, "central_liquidations"),
        "central_replenishments": _len_list(payload, "central_replenishments"),
        "inventory_transfer_proposals": _len_list(payload, "inventory_transfer_proposals"),
        "offer_matches": _len_list(payload, "offer_matches"),
    }


def warehouse_action_stats(payload: dict[str, Any] | None) -> dict[str, int]:
    payload = payload or {}
    stats = {
        "warehouses_controlled": 1 if payload else 0,
        "order_decisions": 0,
        "inventory_offers": 0,
        "inventory_requests": 0,
        "transfer_responses": 0,
        "local_priority": 0,
    }
    if isinstance(payload, dict):
        if "warehouse_actions" in payload and isinstance(payload.get("warehouse_actions"), dict):
            stats["warehouses_controlled"] = len(payload["warehouse_actions"])
            for action in payload["warehouse_actions"].values():
                if isinstance(action, dict):
                    for key in ("order_decisions", "inventory_offers", "inventory_requests", "transfer_responses", "local_priority"):
                        stats[key] += _len_list(action, key)
        else:
            for key in ("order_decisions", "inventory_offers", "inventory_requests", "transfer_responses", "local_priority"):
                stats[key] += _len_list(payload, key)
    return stats


def action_stats(role: str, payload: dict[str, Any] | None) -> dict[str, int]:
    return center_action_stats(payload) if role == "center" else warehouse_action_stats(payload)


def load_model(model_id: str, adapter_id: str | None = None, load_in_4bit: bool = False) -> tuple[Any, Any]:
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    quantization_config = None
    if load_in_4bit:
        compute_dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=compute_dtype,
        )
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        torch_dtype="auto",
        device_map="auto",
        quantization_config=quantization_config,
    )
    if adapter_id:
        log("applying_adapter", adapter_id=adapter_id, load_in_4bit=load_in_4bit)
        model = PeftModel.from_pretrained(model, adapter_id)
    model.eval()
    return model, tokenizer


def rollout(role: str, model: Any, tokenizer: Any, task_id: str, seed: int, max_new_tokens: int) -> dict[str, Any]:
    from supplymind_env_v2.environment import V2SupplyMindEnv
    from supplymind_env_v2.models import CenterAction, V2JointAction, WarehouseAction
    from supplymind_env_v2.policies import fixed_center_action, fixed_warehouse_actions

    env = V2SupplyMindEnv(default_task_id=task_id)
    observation = env.reset_internal(task_id, seed)
    invalid_payloads = 0
    invalid_actions = 0
    steps = 0
    action_totals: Counter[str] = Counter()
    parsed_payloads = 0
    fallback_steps = 0
    samples: list[dict[str, Any]] = []
    while not env.done:
        try:
            if role == "center":
                payload, generated = generate_action(model, tokenizer, role, observation, max_new_tokens)
                parsed_before_validation = payload is not None
                if parsed_before_validation:
                    parsed_payloads += 1
                    action_totals.update(action_stats(role, payload))
                if payload is None:
                    raise ValueError("missing_json")
                action = V2JointAction(
                    warehouse_actions=fixed_warehouse_actions(observation),
                    central_action=CenterAction.model_validate(payload),
                )
            else:
                generated_parts = []
                warehouse_actions = {}
                parsed_before_validation = True
                for warehouse_id in observation.warehouses:
                    payload, generated = generate_action(model, tokenizer, role, observation, max_new_tokens, warehouse_id)
                    generated_parts.append(f"{warehouse_id}: {generated}")
                    if payload is None:
                        parsed_before_validation = False
                        raise ValueError("missing_json")
                    parsed_payloads += 1
                    action_totals.update(action_stats(role, payload))
                    warehouse_actions[warehouse_id] = WarehouseAction.model_validate(payload)
                generated = "\n".join(generated_parts)
                payload = {"warehouse_actions": {key: value.model_dump(mode="json") for key, value in warehouse_actions.items()}}
                action = V2JointAction(warehouse_actions=warehouse_actions, central_action=fixed_center_action(observation, warehouse_actions))
        except Exception:
            invalid_payloads += 1
            fallback_steps += 1
            generated = locals().get("generated", "")
            payload = locals().get("payload", None)
            parsed_before_validation = payload is not None
            if role == "center":
                action = V2JointAction(warehouse_actions=fixed_warehouse_actions(observation), central_action={})
            else:
                action = V2JointAction(warehouse_actions={}, central_action=fixed_center_action(observation, {}))
        result = env.step(action)
        if len(samples) < 3:
            samples.append(
                {
                    "round_index": observation.round_index,
                    "parsed": parsed_before_validation,
                    "fallback": payload is None,
                    "action_stats": action_stats(role, payload),
                    "generated_preview": generated[:600],
                    "parsed_payload": payload,
                    "invalid_action_details": result.observation.feedback.get("invalid_action_details", [])[:5],
                }
            )
        invalid_actions += len(result.observation.feedback.get("invalid_action_details", []))
        observation = result.observation
        steps += 1
    summary = dict(env.last_episode_summary or {})
    summary.update(
        {
            "seed": seed,
            "steps": steps,
            "invalid_payloads": invalid_payloads,
            "invalid_actions": invalid_actions,
            "parsed_payloads": parsed_payloads,
            "fallback_steps": fallback_steps,
            "action_totals": dict(action_totals),
            "samples": samples,
        }
    )
    return summary


def evaluate(role: str, label: str, model: Any, tokenizer: Any, task_id: str, seeds: list[int], max_new_tokens: int) -> dict[str, Any]:
    episodes = [rollout(role, model, tokenizer, task_id, seed, max_new_tokens) for seed in seeds]
    aggregate = {
        "label": label,
        "role": role,
        "episodes": episodes,
        "mean_global_score": mean(float(row.get("graded_score", 0.0)) for row in episodes),
        "mean_center_role_score": mean(float(row.get("center_role_score", 0.0)) for row in episodes),
        "mean_warehouse_role_score": mean(float(row.get("warehouse_role_score", 0.0)) for row in episodes),
        "mean_raw_reward": mean(float(row.get("raw_reward", 0.0)) for row in episodes),
        "invalid_payloads": sum(int(row["invalid_payloads"]) for row in episodes),
        "invalid_actions": sum(int(row["invalid_actions"]) for row in episodes),
        "parsed_payloads": sum(int(row.get("parsed_payloads", 0)) for row in episodes),
        "fallback_steps": sum(int(row.get("fallback_steps", 0)) for row in episodes),
        "action_totals": dict(sum((Counter(row.get("action_totals", {})) for row in episodes), Counter())),
        "sample_generations": [
            {"seed": row.get("seed"), **sample}
            for row in episodes
            for sample in row.get("samples", [])[:1]
        ][:5],
    }
    log("eval_result", **aggregate)
    return aggregate


def main() -> None:
    started = time.time()
    args = parse_args()
    if not (args.adapter_id or args.sft_adapter_id or args.grpo_adapter_id):
        raise SystemExit("Provide --adapter-id or --sft-adapter-id/--grpo-adapter-id.")
    prepare_repo()
    seeds = [int(value.strip()) for value in args.seeds.split(",") if value.strip()]

    log("loading_base_model")
    base_model, tokenizer = load_model(args.model_id, load_in_4bit=args.load_in_4bit)
    base = evaluate(args.role, "base", base_model, tokenizer, args.task_id, seeds, args.max_new_tokens)
    del base_model
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    evaluations: dict[str, Any] = {"base": base}
    adapter_specs = []
    if args.adapter_id:
        adapter_specs.append(("adapter", args.adapter_id))
    if args.sft_adapter_id:
        adapter_specs.append(("sft", args.sft_adapter_id))
    if args.grpo_adapter_id:
        adapter_specs.append(("grpo", args.grpo_adapter_id))
    for label, adapter_id in adapter_specs:
        log("loading_adapter_model", label=label, adapter_id=adapter_id)
        adapter_model, tokenizer = load_model(args.model_id, adapter_id, load_in_4bit=args.load_in_4bit)
        evaluations[label] = evaluate(args.role, label, adapter_model, tokenizer, args.task_id, seeds, args.max_new_tokens)
        del adapter_model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    result = {"message": "eval_done", "elapsed_seconds": round(time.time() - started, 2), **evaluations}
    log(**result)


if __name__ == "__main__":
    main()