bit-jev 0.8.7__py3-none-any.whl

bit_jev/__init__.py

@@ -0,0 +1,18 @@
+"""BitNet 结构化决策：训练、I2_S 导出与逐题 CPU 推理共用同一编码契约。"""
+
+# 包版本与 pyproject.toml、项目更新日志保持同步。
+__version__ = "0.8.7"
+
+# 默认骨干用于新训练运行；蒸馏和 CPU 推理从检查点读取实际来源。
+DEFAULT_BASE = "microsoft/bitnet-b1.58-2B-4T-bf16"
+
+# 顶层 API 延迟导入，单纯查询包版本时不会初始化 tokenizer 或原生构建器。
+__all__ = ["BitJev", "download_model", "__version__"]
+
+
+def __getattr__(name):
+    """按需公开 GGUF 常驻模型与模型下载入口。"""
+    if name in {"BitJev", "download_model"}:
+        from .gguf import BitJev, download_model
+        return {"BitJev": BitJev, "download_model": download_model}[name]
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

bit_jev/__version__.py

@@ -0,0 +1,2 @@
+# 源码包版本与 pyproject.toml 中的发行版本保持一致。
+__version__ = "0.8.7"

bit_jev/_native/CMakeLists.txt

@@ -0,0 +1,26 @@
+cmake_minimum_required(VERSION 3.28)
+project(bit_jev_native LANGUAGES C CXX)
+
+# 只编译本项目所需的 GGUF 推理库与命令行程序；后端由调用方选择。
+set(LLAMA_BUILD_TOOLS OFF CACHE BOOL "" FORCE)
+set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE)
+set(LLAMA_BUILD_TESTS OFF CACHE BOOL "" FORCE)
+set(BUILD_SHARED_LIBS OFF CACHE BOOL "" FORCE)
+
+# 源码构建默认读取工作区依赖；pip 缓存构建显式传入固定版本的路径。
+if(NOT DEFINED BITNET_CPP_DIR)
+    set(BITNET_CPP_DIR "${CMAKE_CURRENT_LIST_DIR}/../../learning/bitnet/3rdparty/llama.cpp")
+endif()
+add_subdirectory("${BITNET_CPP_DIR}" "${CMAKE_BINARY_DIR}/bitnet_cpp")
+
+# 原生程序读取已编码题目与 float32 指针头，主干推理由 llama 库承担。
+add_executable(bit-jev-cpu main.cpp)
+target_compile_features(bit-jev-cpu PRIVATE cxx_std_17)
+# 设备请求必须与编译出的 GPU 后端一致，避免外部 binary 参数静默换后端。
+if(GGML_CUDA)
+    target_compile_definitions(bit-jev-cpu PRIVATE BIT_JEV_BUILT_CUDA=1)
+elseif(GGML_VULKAN)
+    target_compile_definitions(bit-jev-cpu PRIVATE BIT_JEV_BUILT_VULKAN=1)
+endif()
+target_include_directories(bit-jev-cpu PRIVATE "${BITNET_CPP_DIR}/include" "${BITNET_CPP_DIR}/vendor")
+target_link_libraries(bit-jev-cpu PRIVATE llama)

bit_jev/_native/llama-relu2.patch

@@ -0,0 +1,13 @@
+diff --git a/src/models/bitnet.cpp b/src/models/bitnet.cpp
+--- a/src/models/bitnet.cpp
++++ b/src/models/bitnet.cpp
+@@ -130,7 +130,8 @@ llama_model_bitnet::graph::graph(const llama_model & model, const llm_graph_para
+                 model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
+                 NULL,                      NULL, NULL,
+                 NULL,
+-                LLM_FFN_SILU, LLM_FFN_PAR, il);
++                // 蒸馏检查点的 hidden_act 为 relu2，必须与 Transformers 前向保持一致。
++                LLM_FFN_RELU_SQR, LLM_FFN_PAR, il);
+         cb(cur, "ffn_sub_out", il);
+ 
+         cur = build_norm(cur,

bit_jev/_native/main.cpp

@@ -0,0 +1,283 @@
+// bit-jev 原生 GGUF 推理：逐题运行 BitNet，并在候选边界读取隐藏状态。
+#include "llama.h"
+#include "ggml-backend.h"
+#include "nlohmann/json.hpp"
+
+#include <algorithm>
+#include <chrono>
+#include <cmath>
+#include <cstdint>
+#include <cstring>
+#include <fstream>
+#include <iostream>
+#include <memory>
+#include <stdexcept>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+using Json = nlohmann::json;
+
+// 模型加载时仅输出警告和错误，避免每次启动打印整份张量目录。
+void native_log(ggml_log_level level, const char * message, void *) {
+    if (level >= GGML_LOG_LEVEL_WARN) std::cerr << message;
+}
+
+// 指针头文件的矩阵均按行优先排列，顺序与 Python 导出器一致。
+struct Head {
+    uint32_t hidden_size = 0;
+    uint32_t head_size = 0;
+    float temperature = 1.0f;
+    std::vector<float> q_weight;
+    std::vector<float> q_bias;
+    std::vector<float> k_weight;
+    std::vector<float> k_bias;
+};
+
+// llama_batch 持有 C API 分配的内存，在异常路径也自动释放。
+struct BatchHolder {
+    explicit BatchHolder(int capacity) : batch(llama_batch_init(capacity, 0, 1)) {}
+    ~BatchHolder() { llama_batch_free(batch); }
+    llama_batch batch;
+};
+
+// 从固定布局的 sidecar 读取全部指针头参数。
+Head load_head(const std::string & path) {
+    std::ifstream source(path, std::ios::binary);
+    if (!source) throw std::runtime_error("无法打开 head.f32：" + path);
+    char magic[8] = {};
+    Head head;
+    source.read(magic, 8);
+    source.read(reinterpret_cast<char *>(&head.hidden_size), sizeof(head.hidden_size));
+    source.read(reinterpret_cast<char *>(&head.head_size), sizeof(head.head_size));
+    source.read(reinterpret_cast<char *>(&head.temperature), sizeof(head.temperature));
+    if (!source || std::memcmp(magic, "BJHEAD01", 8) != 0 || head.hidden_size == 0 ||
+        head.head_size == 0 || !std::isfinite(head.temperature) || head.temperature <= 0) {
+        throw std::runtime_error("head.f32 文件头无效");
+    }
+    // 四段数组的尺寸由文件头决定，读取长度不足即视为损坏。
+    const size_t matrix_size = static_cast<size_t>(head.hidden_size) * head.head_size;
+    auto read_array = [&source](size_t count) {
+        std::vector<float> values(count);
+        source.read(reinterpret_cast<char *>(values.data()), count * sizeof(float));
+        if (!source) throw std::runtime_error("head.f32 参数数据截断");
+        return values;
+    };
+    head.q_weight = read_array(matrix_size);
+    head.q_bias = read_array(head.head_size);
+    head.k_weight = read_array(matrix_size);
+    head.k_bias = read_array(head.head_size);
+    if (source.peek() != std::char_traits<char>::eof()) {
+        throw std::runtime_error("head.f32 含有多余数据");
+    }
+    return head;
+}
+
+// 将一个骨干隐藏状态投影到指针空间。
+std::vector<float> project(const std::vector<float> & hidden,
+                           const std::vector<float> & weights,
+                           const std::vector<float> & bias,
+                           const Head & head) {
+    std::vector<float> result(head.head_size);
+    for (size_t row = 0; row < head.head_size; ++row) {
+        float value = bias[row];
+        const size_t offset = row * head.hidden_size;
+        for (size_t col = 0; col < head.hidden_size; ++col) {
+            value += weights[offset + col] * hidden[col];
+        }
+        result[row] = value;
+    }
+    return result;
+}
+
+// 清空上一题的 KV 缓存后，按普通因果顺序运行一行。
+std::vector<std::vector<float>> read_hidden(llama_context * context, const Json & row,
+                                             uint32_t hidden_size, int batch_size) {
+    const std::vector<llama_token> ids = row.at("ids").get<std::vector<llama_token>>();
+    const std::vector<int> option_positions = row.at("options").get<std::vector<int>>();
+    const int decide_position = row.at("decide").get<int>();
+    if (ids.empty() || option_positions.empty() || decide_position != static_cast<int>(ids.size()) - 1 ||
+        ids.size() > 4096) {
+        throw std::runtime_error("题目 token 序列或 decide 位置无效");
+    }
+    // 每个输出位置都应位于当前因果行内且互不重复。
+    std::unordered_set<int> wanted(option_positions.begin(), option_positions.end());
+    wanted.insert(decide_position);
+    if (wanted.size() != option_positions.size() + 1) {
+        throw std::runtime_error("候选结束位置重复或与 decide 重叠");
+    }
+    for (int position : wanted) {
+        if (position < 0 || position >= static_cast<int>(ids.size())) {
+            throw std::runtime_error("隐藏状态读取位置越界");
+        }
+    }
+    llama_memory_clear(llama_get_memory(context), true);
+    BatchHolder holder(batch_size);
+    std::vector<std::vector<float>> hidden(ids.size());
+    // 分块送入模型；每次 decode 后立即复制本块要求的隐藏状态。
+    for (size_t start = 0; start < ids.size(); start += batch_size) {
+        const int count = static_cast<int>(std::min<size_t>(batch_size, ids.size() - start));
+        holder.batch.n_tokens = count;
+        for (int index = 0; index < count; ++index) {
+            const int position = static_cast<int>(start) + index;
+            holder.batch.token[index] = ids[position];
+            holder.batch.pos[index] = position;
+            holder.batch.n_seq_id[index] = 1;
+            holder.batch.seq_id[index][0] = 0;
+            // embeddings 模式要求当前块每个 token 都为输出；后续仅复制需要的位置。
+            holder.batch.logits[index] = 1;
+        }
+        if (llama_decode(context, holder.batch) != 0) {
+            throw std::runtime_error("BitNet GGUF decode 失败");
+        }
+        for (int index = 0; index < count; ++index) {
+            const int position = static_cast<int>(start) + index;
+            if (!wanted.count(position)) continue;
+            const float * embedding = llama_get_embeddings_ith(context, index);
+            if (embedding == nullptr) throw std::runtime_error("BitNet 未返回请求的隐藏状态");
+            hidden[position].assign(embedding, embedding + hidden_size);
+        }
+    }
+    // 结果按候选顺序排列，最后一项是 decide 状态。
+    std::vector<std::vector<float>> selected;
+    selected.reserve(option_positions.size() + 1);
+    for (int position : option_positions) selected.push_back(std::move(hidden[position]));
+    selected.push_back(std::move(hidden[decide_position]));
+    return selected;
+}
+
+// 对候选状态执行与 Python PointerHead 完全相同的投影和温度缩放。
+std::vector<float> score_row(const std::vector<std::vector<float>> & states, const Head & head) {
+    const std::vector<float> query = project(states.back(), head.q_weight, head.q_bias, head);
+    const float scale = 1.0f / std::sqrt(static_cast<float>(head.head_size)) / head.temperature;
+    std::vector<float> logits;
+    logits.reserve(states.size() - 1);
+    for (size_t option = 0; option + 1 < states.size(); ++option) {
+        const std::vector<float> key = project(states[option], head.k_weight, head.k_bias, head);
+        float dot = 0.0f;
+        for (size_t index = 0; index < head.head_size; ++index) dot += query[index] * key[index];
+        logits.push_back(dot * scale);
+    }
+    return logits;
+}
+
+// 使用稳定的 softmax 将一题的 logits 转换为概率。
+std::vector<float> softmax(const std::vector<float> & logits) {
+    const float maximum = *std::max_element(logits.begin(), logits.end());
+    std::vector<float> probabilities;
+    probabilities.reserve(logits.size());
+    float sum = 0.0f;
+    for (float value : logits) {
+        const float probability = std::exp(value - maximum);
+        probabilities.push_back(probability);
+        sum += probability;
+    }
+    for (float & probability : probabilities) probability /= sum;
+    return probabilities;
+}
+
+// 读取命令行选项；所有输入均为显式路径或正整数。
+std::string option_value(int argc, char ** argv, const std::string & name,
+                         const std::string & fallback = "") {
+    for (int index = 1; index + 1 < argc; ++index) {
+        if (argv[index] == name) return argv[index + 1];
+    }
+    return fallback;
+}
+
+// 仅检查布尔开关是否出现，不读取其后的路径或数字参数。
+bool has_option(int argc, char ** argv, const std::string & name) {
+    for (int index = 1; index < argc; ++index) {
+        if (argv[index] == name) return true;
+    }
+    return false;
+}
+
+int main(int argc, char ** argv) {
+    const std::string model_path = option_value(argc, argv, "--model");
+    const std::string head_path = option_value(argc, argv, "--head");
+    const std::string device = option_value(argc, argv, "--device", "cpu");
+    const int threads = std::stoi(option_value(argc, argv, "--threads", "4"));
+    const int batch_size = std::stoi(option_value(argc, argv, "--batch", "256"));
+    if (model_path.empty() || head_path.empty() || threads <= 0 || batch_size <= 0 ||
+        (device != "cpu" && device != "vulkan" && device != "cuda")) {
+        std::cerr << "用法：bit-jev-cpu --model 文件.gguf --head head.f32 --device cpu|vulkan|cuda --threads 4 --batch 256\n";
+        return 2;
+    }
+    // 外部传入二进制时也强制匹配后端，不能把 Vulkan 程序当 CUDA 程序运行。
+#if !defined(BIT_JEV_BUILT_VULKAN)
+    if (device == "vulkan") {
+        std::cerr << "本程序未编入 Vulkan 后端\n";
+        return 2;
+    }
+#endif
+#if !defined(BIT_JEV_BUILT_CUDA)
+    if (device == "cuda") {
+        std::cerr << "本程序未编入 CUDA 后端\n";
+        return 2;
+    }
+#endif
+    llama_log_set(native_log, nullptr);
+    llama_backend_init();
+    try {
+        // CPU 不卸载层；GPU 模式申请全部层，实际后端由对应的编译选项提供。
+        llama_model_params model_params = llama_model_default_params();
+        model_params.n_gpu_layers = device == "cpu" ? 0 : 999;
+        if (device != "cpu") {
+            // 同时检查后端能力和可见 GPU，避免用户要求 GPU 时悄悄在 CPU 上运行。
+            if (!llama_supports_gpu_offload() ||
+                ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU) == nullptr) {
+                throw std::runtime_error("当前原生程序没有可用 GPU 后端或可见 GPU；请检查构建选项和驱动");
+            }
+        }
+        std::unique_ptr<llama_model, decltype(&llama_model_free)> model(
+            llama_model_load_from_file(model_path.c_str(), model_params), &llama_model_free);
+        if (!model) throw std::runtime_error("无法加载 I2_S GGUF 模型");
+        const Head head = load_head(head_path);
+        if (head.hidden_size != static_cast<uint32_t>(llama_model_n_embd(model.get()))) {
+            throw std::runtime_error("指针头与 BitNet 隐藏维度不一致");
+        }
+        llama_context_params context_params = llama_context_default_params();
+        context_params.n_ctx = 4096;
+        context_params.n_batch = batch_size;
+        context_params.n_ubatch = batch_size;
+        context_params.n_outputs_max = batch_size;
+        context_params.n_threads = threads;
+        context_params.n_threads_batch = threads;
+        context_params.embeddings = true;
+        context_params.pooling_type = LLAMA_POOLING_TYPE_NONE;
+        std::unique_ptr<llama_context, decltype(&llama_free)> context(
+            llama_init_from_model(model.get(), context_params), &llama_free);
+        if (!context) throw std::runtime_error("无法初始化 BitNet GGUF 上下文");
+        llama_set_embeddings(context.get(), true);
+        // Python API 用握手确认模型已经进入内存或显存，再开始计时和请求。
+        if (has_option(argc, argv, "--ready")) {
+            std::cout << "{\"ready\":true}\n" << std::flush;
+        }
+        // 模型常驻内存；stdin 每一行独立对应 stdout 的一行结果。
+        std::string line;
+        while (std::getline(std::cin, line)) {
+            if (line.empty()) continue;
+            const auto started = std::chrono::steady_clock::now();
+            const Json request = Json::parse(line);
+            Json result;
+            result["logits"] = Json::array();
+            result["probabilities"] = Json::array();
+            for (const Json & row : request.at("rows")) {
+                auto states = read_hidden(context.get(), row, head.hidden_size, batch_size);
+                auto logits = score_row(states, head);
+                result["logits"].push_back(logits);
+                result["probabilities"].push_back(softmax(logits));
+            }
+            const auto ended = std::chrono::steady_clock::now();
+            result["latency_ms"] = std::chrono::duration<double, std::milli>(ended - started).count();
+            std::cout << result.dump() << '\n' << std::flush;
+        }
+    } catch (const std::exception & error) {
+        std::cerr << "bit-jev GGUF 推理失败：" << error.what() << '\n';
+        llama_backend_free();
+        return 1;
+    }
+    llama_backend_free();
+    return 0;
+}

bit_jev/api.py

@@ -0,0 +1,175 @@
+"""TypeSafe-compatible request/response shapes (same contract as kev.api).
+
+Noul   -> 2 options [false, true];             answer = p(true)
+Choice -> options 'name' or 'name: desc';      answer = argmax + probabilities + confidence
+Score  -> options = ordered level descriptions; answer = expected level + probabilities + confidence
+"""
+import json
+from typing import Any, Literal, Union
+
+JSONContent = Union[str, dict, list, int, float, bool, None]
+MAX_OPTIONS = 255
+QUESTION_TYPES = ("noul", "choice", "score")
+
+
+def render(v: JSONContent, indent: int = 0) -> str:
+    """Flatten str | object | array into text the model sees. Field names are kept as labels."""
+    pad = "  " * indent
+    if v is None:
+        return ""
+    if isinstance(v, (str, int, float, bool)):
+        return str(v)
+    if isinstance(v, list):
+        return "\n".join(f"{pad}- {render(x, indent + 1).lstrip()}" for x in v)
+    return "\n".join(
+        f"{pad}{k}:\n{render(x, indent + 1)}" if isinstance(x, (dict, list)) else f"{pad}{k}: {render(x)}"
+        for k, x in v.items()
+    )
+
+
+def option_text(name: str, desc: JSONContent) -> str:
+    return name if desc is None or desc == "" else f"{name}: {render(desc)}"
+
+
+def question_keys(qtype: str, criteria) -> list:
+    """The keys a question's probabilities are reported under, in option order."""
+    if qtype == "choice":
+        return list(criteria)
+    if qtype == "noul":
+        return ["false", "true"]
+    return [str(i) for i in range(len(criteria))]
+
+
+def validate_request(req: dict) -> None:
+    """Raise ValueError on a malformed SystemOne-shaped request dict."""
+    if not isinstance(req, dict) or "questions" not in req:
+        raise ValueError("request needs a 'questions' object")
+    if not req["questions"]:
+        raise ValueError("at least one question is required")
+    for qid, q in req["questions"].items():
+        t = q.get("type")
+        if t not in QUESTION_TYPES:
+            raise ValueError(f"question {qid!r}: type must be one of {QUESTION_TYPES}")
+        if t in ("choice", "score"):
+            crit = q.get("criteria")
+            n = len(crit) if crit else 0
+            if not 1 <= n <= MAX_OPTIONS:
+                raise ValueError(f"question {qid!r}: criteria must have 1..{MAX_OPTIONS} options")
+
+
+def to_record(req: dict, labelled: bool = True):
+    """SystemOne-shaped dict -> (internal record, per-question meta).
+
+    With labelled=False, labels are zero placeholders and may be absent."""
+    validate_request(req)
+    qs, meta = [], []
+    for qid, q in req["questions"].items():
+        t = q["type"]
+        m = {"id": qid, "type": t, "keys": question_keys(t, q.get("criteria"))}
+        if t == "noul":
+            c = q.get("criteria") or {}
+            opts = [option_text("no", c.get("false")), option_text("yes", c.get("true"))]
+            label = int(q["label"]) if "label" in q else -1
+        elif t == "choice":
+            opts = [option_text(k, v) for k, v in q["criteria"].items()]
+            label = m["keys"].index(q["label"]) if "label" in q else -1
+        else:
+            opts = [render(x) for x in q["criteria"]]
+            m["legend"] = dict(zip(m["keys"], opts))
+            label = int(q["label"]) if "label" in q else -1
+        if labelled and label < 0:
+            raise ValueError(f"question {qid!r} of type {t} is missing a label")
+        qs.append({"instr": render(q.get("instructions")), "options": opts,
+                   "label": max(label, 0), "qtype": t})
+        meta.append(m)
+    return {"state": render(req.get("state")), "questions": qs}, meta
+
+
+def load_requests(path):
+    """Labelled requests from a JSONL file: one SystemOne-shaped object per line, a label per question."""
+    reqs = []
+    with open(path, encoding="utf-8") as f:
+        for i, line in enumerate(f, 1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                r = json.loads(line)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"{path}:{i}: invalid JSON: {e}") from e
+            r.setdefault("_meta", {})
+            r["_meta"].setdefault("id", f"{path}:{i}")
+            r["_meta"].setdefault("source", "user")
+            reqs.append(r)
+    if not reqs:
+        raise ValueError(f"{path}: no records found")
+    return reqs
+
+
+def read_jsonl(path):
+    with open(path, encoding="utf-8") as f:
+        return [json.loads(line) for line in f if line.strip()]
+
+
+def write_jsonl(path, rows):
+    with open(path, "w", encoding="utf-8", newline="\n") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+
+
+def write_json(path, obj):
+    with open(path, "w", encoding="utf-8", newline="\n") as f:
+        json.dump(obj, f, ensure_ascii=False, indent=2)
+
+
+def read_json(path):
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def choice_confidence(p) -> float:
+    K = len(p)
+    return 1.0 if K == 1 else (max(p) - 1 / K) / (1 - 1 / K)
+
+
+def score_confidence(p) -> float:
+    """1 - E|level - mode| / (L - 1): how concentrated the distribution is at its modal level."""
+    L = len(p)
+    if L == 1:
+        return 1.0
+    mode = max(range(L), key=lambda i: p[i])
+    return 1.0 - sum(pi * abs(i - mode) for i, pi in enumerate(p)) / (L - 1)
+
+
+def round_prob(x: float) -> float:
+    return round(float(x), 4)
+
+
+def to_answers(probs, meta) -> dict:
+    """Per-question probability lists -> TypeSafe-shaped answers (kev.api.to_answers)."""
+    out = {}
+    for p, m in zip(probs, meta):
+        p = [float(x) for x in p]
+        if m["type"] == "noul":
+            out[m["id"]] = {"type": "noul", "noul": round_prob(p[1])}
+        elif m["type"] == "choice":
+            out[m["id"]] = {
+                "type": "choice",
+                "choice": m["keys"][max(range(len(p)), key=lambda i: p[i])],
+                "confidence": round_prob(choice_confidence(p)),
+                "probabilities": {k: round_prob(v) for k, v in zip(m["keys"], p)},
+            }
+        else:
+            out[m["id"]] = {
+                "type": "score",
+                "score": round_prob(sum(i * pi for i, pi in enumerate(p))),
+                "confidence": round_prob(score_confidence(p)),
+                "legend": m["legend"],
+                "probabilities": {str(i): round_prob(v) for i, v in enumerate(p)},
+            }
+    return out
+
+
+def output_tokens(tok, answers: dict) -> int:
+    """Billing-style figure: tokens of the serialized answers (there is no generation)."""
+    return len(tok(json.dumps(answers), add_special_tokens=False).input_ids)

bit_jev/checkpoint.py

@@ -0,0 +1,149 @@
+"""Trained checkpoints: a directory (or Hub repo) holding a LoRA adapter, `head.pt` and the tokenizer.
+
+    ck = BitJevCheckpoint("runs/bit-jev-2b")        # or a hub id
+    tok, model = ck.load("cuda", dtype=torch.float32)
+
+`head.pt` carries the architecture (Meta) and the pointer-head state dict; the load path is the only
+place that knows both, so train / eval / serve / convert all build the same model.
+"""
+import os
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import torch
+
+from . import hub as _hub
+from .model import DecisionModel, load_tokenizer
+
+HUB_ID = re.compile(r"[\w.-]+/[\w.-]+(@[\w.-]+)?")
+
+
+def is_hub_id(run):
+    return not os.path.isdir(run) and HUB_ID.fullmatch(str(run)) is not None
+
+
+def resolve_run(run):
+    """Local directory as given, or a Hub repo id (optionally @revision) downloaded to the HF cache."""
+    if os.path.isdir(run):
+        return str(run)
+    repo, _, revision = str(run).partition("@")
+    return _hub.snapshot(repo, revision=revision or None,
+                         allow_patterns=["*.json", "*.safetensors", "*.pt", "*.txt", "*.jinja", "*.model"])
+
+
+@dataclass
+class Meta:
+    """Contents of head.pt. `extra` keeps everything else in the file (training args, hashes, fits)."""
+    base: str
+    head: dict | None = None
+    base_revision: str | None = None
+    lora: int = 0
+    head_dim: int = 256
+    temperature: float = 1.0
+    arch: str = "bit-jev/0.1"
+    holdout: list = field(default_factory=list)
+    extra: dict = field(default_factory=dict)
+
+    KNOWN = ("base", "head", "base_revision", "lora", "head_dim", "temperature", "arch", "holdout")
+
+    @classmethod
+    def from_dict(cls, d):
+        if d.get("arch") not in (None, "bit-jev/0.1"):
+            raise ValueError(f"head.pt arch {d.get('arch')!r} is not a bit-jev checkpoint")
+        return cls(**{k: d[k] for k in cls.KNOWN if k in d},
+                   extra={k: v for k, v in d.items() if k not in cls.KNOWN})
+
+    def to_dict(self):
+        return {**self.extra, **{k: getattr(self, k) for k in self.KNOWN}}
+
+
+def read_meta(run):
+    return Meta.from_dict(torch.load(f"{run}/head.pt", map_location="cpu", weights_only=False))
+
+
+def write_meta(run, meta):
+    torch.save(meta.to_dict(), f"{run}/head.pt")
+
+
+class BitJevCheckpoint:
+    def __init__(self, run):
+        self.requested = str(run)
+        self.path = resolve_run(run)
+        self.meta = read_meta(self.path)
+
+    def file(self, name):
+        return Path(self.path) / name
+
+    def load(self, device, dtype=torch.float32, merge=None, temperature=None):
+        """-> (tokenizer, model) in eval mode.
+
+        merge=None defaults by base: on a BitNet base the LoRA adapter must stay UNMERGED. The
+        online quantizer (transformers' bitnet integration) rounds every weight to {-1,0,+1} at
+        each forward, so a merged delta is quantized away and the served model silently loses
+        everything the adapter learned (73.1% unmerged vs 32.4% merged on decision-v7 dev).
+        An unmerged LoRA side-branch computes in full precision and reproduces training exactly.
+        merge=True remains available for dense fp bases (kev-style), where merging is exact.
+
+        Distilled runs (bit_jev.distill) store a full backbone instead of an adapter: a run dir
+        with config.json and no adapter_model.safetensors is loaded as the backbone itself."""
+        has_adapter = self.file("adapter_model.safetensors").exists()
+        meta = self.meta
+        if has_adapter or not self.file("tokenizer_config.json").exists():
+            tok = load_tokenizer(meta.base, revision=meta.base_revision)
+        else:
+            tok = load_tokenizer(self.path)   # distilled run: the run dir carries its own tokenizer
+        if has_adapter:
+            backbone = meta.base
+        else:
+            if not self.file("config.json").exists():
+                raise FileNotFoundError(f"{self.path} has neither an adapter nor a full backbone "
+                                        "(config.json missing); not a bit-jev run dir")
+            backbone = self.path
+        m = DecisionModel(backbone, tok, device, lora=None,
+                          revision=meta.base_revision if has_adapter else None,
+                          head_dim=meta.head_dim, dtype=dtype)
+        if has_adapter:
+            from peft import PeftModel
+            m.lm = PeftModel.from_pretrained(m.lm, self.path, torch_device=str(device)).to(device)
+            if merge is None:
+                merge = not self._is_bitnet_base()
+            if merge:
+                m.lm = m.lm.merge_and_unload()
+        if dtype != torch.float32:
+            m.lm = m.lm.to(dtype)
+        m.head.load_state_dict(meta.head)
+        m.head.temperature = meta.temperature if temperature is None else float(temperature)
+        m.eval()
+        return tok, m
+
+    def _is_bitnet_base(self):
+        """True when the checkpoint's base uses BitNet's online ternary quantization."""
+        if "bitnet" in (self.meta.base or "").lower():
+            return True
+        try:
+            from . import hub as _hub
+            cfg = _hub.load_config(self.meta.base, revision=self.meta.base_revision)
+            return getattr(cfg, "model_type", "") == "bitnet"
+        except Exception:
+            return False
+
+    COMPAT_FIELDS = ("base", "lora", "head_dim")
+
+    def warm_start(self, model, ours: Meta):
+        """Delta training: load this checkpoint's adapter and head into a fresh DecisionModel;
+        architecture fields are compared before loading (a silent mismatch still "trains" otherwise)."""
+        from peft import get_peft_model_state_dict, load_peft_weights, set_peft_model_state_dict
+        for name in self.COMPAT_FIELDS:
+            if getattr(self.meta, name) != getattr(ours, name):
+                raise ValueError(f"--init_from {self.path}: {name} differs "
+                                 f"({getattr(self.meta, name)!r} vs {getattr(ours, name)!r})")
+        weights = load_peft_weights(self.path, device="cpu")
+        have = set(get_peft_model_state_dict(model.lm))
+        if set(weights) != have:
+            raise ValueError(f"--init_from {self.path}: adapter tensors mismatch "
+                             f"(missing {sorted(have - set(weights))[:2]}, "
+                             f"unexpected {sorted(set(weights) - have)[:2]})")
+        set_peft_model_state_dict(model.lm, weights)
+        model.head.load_state_dict(self.meta.head)
+        return {"init_from": self.requested, "resolved": self.path, "adapter_tensors": len(weights)}