Add cached Hub loading for mini-Jev

README.md

@@ -2,6 +2,8 @@
 license: apache-2.0
 base_model: Qwen/Qwen3-0.6B
 library_name: custom
+datasets:
+  - samatv256/jev-decisions-v1
 tags:
   - mini-jev
   - decision-model
@@ -49,11 +51,19 @@ The current strength is **finite-choice tool and action selection** from state a
 
 ## Use
 
-A CUDA GPU with BF16 support is required for this 4-bit package. Install the tested dependencies from requirements.txt. The [Qwen3-0.6B base model](https://huggingface.co/Qwen/Qwen3-0.6B) is downloaded separately on first load; this repository contains only the adapter and decision head.
+A CUDA GPU with BF16 support is required for this 4-bit package. Install the tested dependencies from requirements.txt. The [Qwen3-0.6B base model](https://huggingface.co/Qwen/Qwen3-0.6B) is downloaded separately on first load; the current checkpoint consists of the adapter and decision-head files described above.
+
+Fetch the inference module and dependency list into a working directory first:
+
+```bash
+python -m pip install "huggingface_hub==1.33.0"
+hf download samatv256/mini-Jev inference.py requirements.txt --local-dir .
+python -m pip install -r requirements.txt
+```
 
     from inference import MiniJev
 
-    model = MiniJev.load(".")
+    model = MiniJev.load()  # Defaults to samatv256/mini-Jev using the HF cache.
     result = model.predict(
         state={"user_goal": "Find tomorrow's weather forecast for Paris."},
         question="Given the current state and available options,\nwhich option should be selected?",
@@ -70,6 +80,15 @@ A CUDA GPU with BF16 support is required for this 4-bit package. Install the tes
     print(result["selected_id"])
     print(result["options"])
 
+To pin a release, use `MiniJev.load("samatv256/mini-Jev", revision="step-010626")`.
+Existing local directories and `Path` arguments remain supported, for example
+`MiniJev.load(".")` after downloading the release files. Existing local paths take
+precedence over repo IDs. `revision` applies only to Hub loading. Set
+`HF_HUB_OFFLINE=1` to use cached files offline; both the release and its pinned Qwen
+base/tokenizer must already be cached. Hub loading fetches only `config.json`,
+`decision_head.safetensors`, `adapter/adapter_config.json`, and
+`adapter/adapter_model.safetensors` from one revision, excluding legacy root weights.
+
 Option IDs are bookkeeping and are excluded from model text. The model scores each supplied option and normalizes probabilities over the finite set. Input branches are limited to 8,192 tokens.
 
 ## Data attribution

inference.py

@@ -7,6 +7,8 @@ from pathlib import Path
 from typing import Any
 
 import torch
+from huggingface_hub import snapshot_download
+from huggingface_hub.utils import HFValidationError, validate_repo_id
 from peft import PeftModel, prepare_model_for_kbit_training
 from safetensors.torch import load_file
 from torch import nn
@@ -215,7 +217,37 @@ class MiniJev:
         self.max_tokens = int(config["max_tokens"])
 
     @classmethod
-    def load(cls, release_dir: str | Path, device: str = "cuda") -> MiniJev:
+    def load(
+        cls,
+        release_dir: str | Path = "samatv256/mini-Jev",
+        device: str = "cuda",
+        *,
+        revision: str | None = None,
+    ) -> MiniJev:
+        """Load a local release or a cached Hub namespace/repo at a revision.
+
+        Existing local paths and Path arguments always remain local. Revision
+        applies only to Hub releases. HF_HUB_OFFLINE uses the normal Hub cache.
+        """
+        if isinstance(release_dir, str) and "/" in release_dir and not Path(release_dir).exists():
+            try:
+                validate_repo_id(release_dir)
+            except HFValidationError:
+                pass  # Explicit or invalid local paths retain constructor behavior.
+            else:
+                files = [
+                    "config.json",
+                    "decision_head.safetensors",
+                    "adapter/adapter_config.json",
+                    "adapter/adapter_model.safetensors",
+                ]
+                root = Path(snapshot_download(
+                    repo_id=release_dir, revision=revision, allow_patterns=files,
+                ))
+                missing = [name for name in files if not (root / name).is_file()]
+                if missing:
+                    raise FileNotFoundError("missing required release files: " + ", ".join(missing))
+                return cls(root, device)
         return cls(release_dir, device)
 
     def predict(

requirements.txt

@@ -4,3 +4,4 @@ peft==0.21.0
 bitsandbytes==0.50.2
 safetensors==0.8.0
 accelerate==1.15.0
+huggingface_hub==1.33.0

tests/test_release_loading.py

@@ -0,0 +1,147 @@
+"""Release-location tests; model construction is replaced to avoid GPU work."""
+
+import json
+import re
+from pathlib import Path
+
+import pytest
+from huggingface_hub import constants
+from huggingface_hub.errors import LocalEntryNotFoundError
+
+import inference
+
+FILES = (
+    "config.json",
+    "decision_head.safetensors",
+    "adapter/adapter_config.json",
+    "adapter/adapter_model.safetensors",
+)
+
+
+class ProbeMiniJev(inference.MiniJev):
+    def __init__(self, release_dir, device="cuda"):
+        self.release_dir = release_dir
+        self.device = device
+        self.config = json.loads((Path(release_dir) / "config.json").read_text())
+
+
+@pytest.fixture
+def release(tmp_path):
+    for name in FILES:
+        path = tmp_path / name
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text('{"marker": "resolved"}' if name.endswith('.json') else 'weights')
+    return tmp_path
+
+
+@pytest.fixture
+def hub(monkeypatch, release):
+    calls = []
+
+    def download(**kwargs):
+        calls.append(kwargs)
+        return str(release)
+
+    monkeypatch.setattr(inference, "snapshot_download", download, raising=False)
+    return calls
+
+
+def test_local_string_takes_precedence_over_repo_id(monkeypatch, tmp_path, release, hub):
+    local = tmp_path / "example" / "model"
+    local.mkdir(parents=True)
+    (local / "config.json").write_text('{"marker": "local"}')
+    monkeypatch.chdir(tmp_path)
+    model = ProbeMiniJev.load("example/model", "cuda:1")
+    assert model.config["marker"] == "local"
+    assert model.release_dir == "example/model"
+    assert model.device == "cuda:1"
+    assert hub == []
+
+
+def test_path_remains_local(release, hub):
+    model = ProbeMiniJev.load(release)
+    assert model.release_dir is release
+    assert model.config["marker"] == "resolved"
+    assert hub == []
+
+
+@pytest.mark.parametrize("argument", [None, "samatv256/mini-Jev", "example/model"])
+def test_hub_default_and_explicit_repo_use_exact_allowlist(argument, release, hub):
+    model = ProbeMiniJev.load() if argument is None else ProbeMiniJev.load(argument)
+    assert model.config["marker"] == "resolved"
+    assert Path(model.release_dir) == release
+    assert len(hub) == 1
+    assert hub[0] == {
+        "repo_id": argument or "samatv256/mini-Jev",
+        "revision": None,
+        "allow_patterns": list(FILES),
+    }
+
+
+@pytest.mark.parametrize("revision", ["step-010626", "a" * 40])
+def test_revision_is_resolved_in_one_snapshot(revision, release, hub):
+    model = ProbeMiniJev.load("example/model", revision=revision)
+    assert model.config["marker"] == "resolved"
+    assert len(hub) == 1
+    assert hub[0]["revision"] == revision
+
+
+@pytest.mark.parametrize("missing", FILES)
+def test_incomplete_hub_snapshot_fails_before_constructor(missing, release, hub):
+    (release / missing).unlink()
+    with pytest.raises(FileNotFoundError, match="missing required release files:.*" + re.escape(missing)):
+        ProbeMiniJev.load("example/model")
+
+
+@pytest.mark.parametrize("argument", ["./missing/model", "../missing/model", "/missing/model", "missing", Path("missing/model")])
+def test_explicit_missing_local_paths_do_not_download(argument, hub):
+    with pytest.raises(FileNotFoundError):
+        ProbeMiniJev.load(argument)
+    assert hub == []
+
+
+def test_local_revision_does_not_change_loading(release, hub):
+    model = ProbeMiniJev.load(str(release), revision="step-010626")
+    assert model.config["marker"] == "resolved"
+    assert hub == []
+
+
+def test_offline_cache_miss_is_propagated(monkeypatch):
+    def download(**kwargs):
+        raise LocalEntryNotFoundError("no cached release")
+
+    monkeypatch.setattr(inference, "snapshot_download", download, raising=False)
+    with pytest.raises(LocalEntryNotFoundError, match="no cached release"):
+        ProbeMiniJev.load("example/model")
+
+
+@pytest.mark.parametrize("revision", [None, "a" * 40])
+def test_real_hub_offline_cache_hit(monkeypatch, tmp_path, release, revision):
+    cache = tmp_path / "cache"
+    storage = cache / "models--example--model"
+    snapshot = storage / "snapshots" / ("a" * 40)
+    entries = {}
+    for name in FILES:
+        data = (release / name).read_bytes()
+        path = snapshot / name
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_bytes(data)
+        entries[name] = {"size": len(data), "blob_id": "b" * 40}
+    (storage / "refs").mkdir()
+    (storage / "refs" / "main").write_text("a" * 40)
+    (storage / "trees").mkdir()
+    (storage / "trees" / (("a" * 40) + ".json")).write_text(
+        json.dumps({"format_version": 1, "files": entries})
+    )
+    monkeypatch.setattr(constants, "HF_HUB_CACHE", str(cache))
+    monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
+    model = ProbeMiniJev.load("example/model", revision=revision)
+    assert Path(model.release_dir) == snapshot
+    assert model.config["marker"] == "resolved"
+
+
+def test_real_hub_offline_cache_miss(monkeypatch, tmp_path):
+    monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path / "empty-cache"))
+    monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
+    with pytest.raises(LocalEntryNotFoundError):
+        ProbeMiniJev.load("example/uncached")