| """Active trained-model discovery for product inference. |
| |
| The HF training Space writes model artifacts into per-sweep folders. The app |
| uses this module to find the locally activated artifact without hard-coding a |
| specific checkpoint path into the API or agent stack. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from pathlib import Path |
| from typing import Any |
|
|
|
|
| ROOT = Path(__file__).resolve().parents[3] |
| ACTIVE_DIR = ROOT / "checkpoints" / "active" |
| MANIFEST_PATH = ACTIVE_DIR / "active_model_manifest.json" |
| DEFAULT_RUN_ID = "qwen-qwen2-5-0-5b-instruct" |
|
|
|
|
| def _truthy(value: str | None) -> bool | None: |
| if value is None: |
| return None |
| lowered = value.strip().lower() |
| if lowered in {"1", "true", "yes", "on"}: |
| return True |
| if lowered in {"0", "false", "no", "off"}: |
| return False |
| return None |
|
|
|
|
| def _read_json(path: Path) -> dict[str, Any]: |
| if not path.exists(): |
| return {} |
| try: |
| payload = json.loads(path.read_text(encoding="utf-8")) |
| except Exception: |
| return {} |
| return payload if isinstance(payload, dict) else {} |
|
|
|
|
| def _resolve_path(value: str | Path | None, default: Path) -> Path: |
| if value is None or str(value).strip() == "": |
| return default |
| path = Path(str(value)).expanduser() |
| if path.is_absolute(): |
| return path |
| return ROOT / path |
|
|
|
|
| def _adapter_base_model(adapter_dir: Path) -> str: |
| payload = _read_json(adapter_dir / "adapter_config.json") |
| value = payload.get("base_model_name_or_path") |
| return str(value) if isinstance(value, str) else "" |
|
|
|
|
| def active_model_status() -> dict[str, Any]: |
| """Return the activated model artifact contract used by the app.""" |
|
|
| manifest = _read_json(MANIFEST_PATH) |
| env_enabled = _truthy(os.getenv("POLYGUARD_ENABLE_ACTIVE_MODEL")) |
| manifest_enabled = bool(manifest.get("enabled", False)) |
| enabled = env_enabled if env_enabled is not None else manifest_enabled |
|
|
| preferred_artifact = ( |
| os.getenv("POLYGUARD_ACTIVE_PREFERRED_ARTIFACT") |
| or str(manifest.get("preferred_artifact") or "grpo_adapter") |
| ) |
| if preferred_artifact not in {"grpo_adapter", "merged", "sft_adapter"}: |
| preferred_artifact = "grpo_adapter" |
|
|
| grpo_adapter = _resolve_path( |
| os.getenv("POLYGUARD_ACTIVE_GRPO_ADAPTER") or manifest.get("grpo_adapter"), |
| ACTIVE_DIR / "grpo_adapter", |
| ) |
| sft_adapter = _resolve_path( |
| os.getenv("POLYGUARD_ACTIVE_SFT_ADAPTER") or manifest.get("sft_adapter"), |
| ACTIVE_DIR / "sft_adapter", |
| ) |
| merged_model = _resolve_path( |
| os.getenv("POLYGUARD_ACTIVE_MERGED_MODEL") or manifest.get("merged_model"), |
| ACTIVE_DIR / "merged", |
| ) |
| base_model = ( |
| os.getenv("POLYGUARD_ACTIVE_BASE_MODEL") |
| or str(manifest.get("base_model") or "") |
| or _adapter_base_model(grpo_adapter) |
| or _adapter_base_model(sft_adapter) |
| or os.getenv("POLYGUARD_HF_MODEL", "Qwen/Qwen2.5-0.5B-Instruct") |
| ) |
|
|
| availability = { |
| "grpo_adapter": (grpo_adapter / "adapter_config.json").exists() |
| and (grpo_adapter / "adapter_model.safetensors").exists(), |
| "merged": (merged_model / "config.json").exists(), |
| "sft_adapter": (sft_adapter / "adapter_config.json").exists() |
| and (sft_adapter / "adapter_model.safetensors").exists(), |
| } |
| load_order = [preferred_artifact] + [ |
| item for item in ["grpo_adapter", "merged", "sft_adapter"] if item != preferred_artifact |
| ] |
| active = any(availability.values()) |
|
|
| return { |
| "enabled": enabled, |
| "active": active, |
| "manifest_path": str(MANIFEST_PATH), |
| "manifest_exists": MANIFEST_PATH.exists(), |
| "run_id": str(manifest.get("run_id") or DEFAULT_RUN_ID), |
| "source": str(manifest.get("source") or ""), |
| "label": str(manifest.get("label") or ""), |
| "model_id": str(manifest.get("model_id") or base_model), |
| "base_model": base_model, |
| "preferred_artifact": preferred_artifact, |
| "load_order": load_order, |
| "availability": availability, |
| "paths": { |
| "grpo_adapter": str(grpo_adapter), |
| "merged": str(merged_model), |
| "sft_adapter": str(sft_adapter), |
| }, |
| "reports": manifest.get("reports", {}) if isinstance(manifest.get("reports"), dict) else {}, |
| "notes": str(manifest.get("notes") or ""), |
| } |
|
|
|
|
| def available_artifact_path(status: dict[str, Any] | None = None) -> tuple[str, Path] | None: |
| """Return the first available artifact according to the active load order.""" |
|
|
| status = status or active_model_status() |
| if not status.get("enabled") or not status.get("active"): |
| return None |
| paths = status.get("paths", {}) |
| availability = status.get("availability", {}) |
| if not isinstance(paths, dict) or not isinstance(availability, dict): |
| return None |
| for artifact in status.get("load_order", []): |
| if availability.get(artifact) and paths.get(artifact): |
| return str(artifact), Path(str(paths[artifact])) |
| return None |
|
|
