Fix adapter eval for 3B notebooks

scripts/hf_eval_supplymind_adapters.py

@@ -25,7 +25,7 @@ from typing import Any
 import torch
 from huggingface_hub import snapshot_download
 from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 
 REPO_ID = "rishavutk/supplymind"
@@ -46,6 +46,7 @@ def parse_args() -> argparse.Namespace:
     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()
 
 
@@ -117,13 +118,19 @@ def compact_observation(observation: Any, role: str, warehouse_id: str | None =
 def system_prompt(role: str) -> str:
     if role == "center":
         return (
-            "You are the center policy in SupplyMind. Return only strict JSON matching CenterAction. "
-            "Use central_procurements, central_liquidations, central_replenishments, inventory_transfer_proposals, offer_matches."
+            "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. You control exactly one warehouse from the user observation. "
-        "Return only strict JSON matching WarehouseAction with keys order_decisions, inventory_offers, inventory_requests, "
-        "transfer_responses, and local_priority. Only use visible order_id and proposal_id values."
+        "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."
     )
 
 
@@ -195,10 +202,25 @@ 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) -> tuple[Any, Any]:
+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)
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
+    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
@@ -324,7 +346,7 @@ def main() -> None:
     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)
+    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():
@@ -340,7 +362,7 @@ def main() -> None:
         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)
+        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():