| import os |
| import torch |
| import logging |
|
|
| logger = logging.getLogger(__name__) |
|
|
| from model import RouterNetwork |
|
|
|
|
| class RouterPolicy: |
| """Thin wrapper around RouterNetwork for loading, inference and checkpointing.""" |
|
|
| def __init__(self, model: RouterNetwork, device: str = "cpu"): |
| self.model = model.to(device) |
| self.device = torch.device(device) |
| self.model.eval() |
|
|
| @classmethod |
| def load(cls, checkpoint_path: str, embedding_dim: int | None = None, device: str = "cpu"): |
| """Load RouterNetwork from checkpoint. If checkpoint contains container dict with |
| 'model_state_dict' will use that, otherwise will try to load state_dict directly. |
| |
| embedding_dim may be required if checkpoint does not include meta info. |
| """ |
| checkpoint_path = os.path.abspath(checkpoint_path) |
| if not os.path.exists(checkpoint_path): |
| raise FileNotFoundError(f"Router checkpoint not found: {checkpoint_path}") |
|
|
| sd = torch.load(checkpoint_path, map_location="cpu") |
|
|
| |
| if isinstance(sd, dict) and "model_state_dict" in sd: |
| model_state = sd["model_state_dict"] |
| meta = sd.get("meta", {}) |
| else: |
| model_state = sd |
| meta = {} |
|
|
| |
| emb_dim = embedding_dim or meta.get("embedding_dim") or meta.get("embed_dim") or 384 |
|
|
| model = RouterNetwork(embedding_dim=emb_dim) |
|
|
| missing, unexpected = model.load_state_dict(model_state, strict=False) |
| if missing or unexpected: |
| logger.info(f"RouterPolicy.load: missing keys={list(missing)} unexpected={list(unexpected)}") |
|
|
| return cls(model, device=device) |
|
|
| def predict_sync(self, emb_tensor: torch.Tensor): |
| """Synchronous prediction: accepts a 1-D or 2-D embedding tensor and returns |
| (action_idx:int, prob:float, value:float). |
| """ |
| was_1d = False |
| if emb_tensor.dim() == 1: |
| emb_tensor = emb_tensor.unsqueeze(0) |
| was_1d = True |
|
|
| emb = emb_tensor.to(self.device) |
|
|
| with torch.no_grad(): |
| logits, value = self.model(emb) |
| probs = torch.softmax(logits, dim=-1) |
| top_idx = int(probs.argmax(dim=-1)[0].item()) |
| top_prob = float(probs[0, top_idx].item()) |
| val = float(value.squeeze(1)[0].item()) if value is not None else 0.0 |
|
|
| return top_idx, top_prob, val |
|
|
| def save_checkpoint(self, path: str, optimizer: torch.optim.Optimizer | None = None, meta: dict | None = None): |
| path = os.path.abspath(path) |
| payload = { |
| "model_state_dict": self.model.state_dict(), |
| "meta": meta or {} |
| } |
| if optimizer is not None: |
| try: |
| payload["optimizer_state_dict"] = optimizer.state_dict() |
| except Exception: |
| logger.warning("Could not include optimizer state in checkpoint") |
|
|
| torch.save(payload, path) |
| return path |
|
|
