7e3eeee048
* model : add H2O-Lightning-4B decision model Adds the h2o-lightning decision type for h2oai/h2o-lightning-4b (Qwen3.5-4B with a merged decision LoRA). None of the existing types match this model. It answers with " A" rather than "A", uses up to 255 labels, and was trained on Python json.dumps text, so floats like 12.0 must render as 12.0. The converter writes the decision metadata and the systemone template from the model's serve_config.json and refuses a config that doesn't match. Tests skip unless the tiny test model is available locally. Assisted-by: Claude Opus 5.5 * rm redundant code * put labels to gguf --------- Co-authored-by: Jonathan C. McKinney <jon.mckinney@h2o.ai> Co-authored-by: Xuan Son Nguyen <son@huggingface.co>
97 lines
4.5 KiB
Python
97 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .base import ModelBase, gguf, logger
|
|
from .qwen import Qwen3_5TextModel
|
|
from .qwen3vl import Qwen3VLVisionModel
|
|
|
|
|
|
def _is_h2o_lightning_checkpoint(dir_model: Path) -> bool:
|
|
return (dir_model / "h2o_lightning_shim.py").is_file() and (dir_model / "serve_config.json").is_file()
|
|
|
|
|
|
@ModelBase.register_hparams_loader(_is_h2o_lightning_checkpoint)
|
|
def _load_h2o_lightning_hparams(dir_model: Path) -> dict[str, Any]:
|
|
logger.info("gguf: detected H2O-Lightning checkpoint")
|
|
hparams = ModelBase.load_hparams(dir_model, False, guess=False)
|
|
hparams["architectures"] = ["H2OLightningModel"]
|
|
return hparams
|
|
|
|
|
|
@ModelBase.register("H2OLightningModel")
|
|
@ModelBase.example("h2oai/h2o-lightning-4b")
|
|
class H2OLightningModel(Qwen3_5TextModel):
|
|
model_arch = gguf.MODEL_ARCH.QWEN35
|
|
|
|
# prompt and calibration follow serve_config.json and h2o_lightning_shim.py of the model repo
|
|
# the state budget is hard-coded in server-decision.cpp
|
|
_MAX_STATE_TOKENS = 32000
|
|
_DEFAULT_INSTRUCTIONS = "Answer the question below."
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
with open(self.dir_model / "serve_config.json", "r", encoding="utf-8") as f:
|
|
self._serve_config = json.load(f)
|
|
prompt = self._serve_config["prompt"]
|
|
expected = {
|
|
"state_compact": (prompt.get("state_compact", True), True),
|
|
"max_state_tokens": (prompt.get("max_state_tokens"), self._MAX_STATE_TOKENS),
|
|
"layout": (prompt.get("layout", "plain"), "plain"),
|
|
"chat_template": (prompt.get("chat_template", True), True),
|
|
"enable_thinking": (prompt.get("enable_thinking", False), False),
|
|
}
|
|
for name, (got, want) in expected.items():
|
|
if got != want:
|
|
raise ValueError(f"serve_config.json: {name} is {got!r}, the h2o-lightning decision type expects {want!r}")
|
|
|
|
def get_vocab_base_pre(self, tokenizer) -> str:
|
|
# transformers loads this checkpoint as Qwen2Tokenizer: the model was trained and is served with the Qwen2 split
|
|
pre = super().get_vocab_base_pre(tokenizer)
|
|
if pre != "qwen35":
|
|
raise ValueError(f"expected the Qwen3.5 tokenizer, got pre-tokenizer {pre!r}")
|
|
return "qwen2"
|
|
|
|
def set_vocab(self):
|
|
super().set_vocab()
|
|
self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
|
|
|
|
def _systemone_template(self) -> str:
|
|
# the server gives the template texts rendered as the shim renders them, it only lays them out
|
|
prompt = self._serve_config["prompt"]
|
|
nl = "{{ '\\n' }}"
|
|
description = (
|
|
"{% if type == 'noul' %}"
|
|
"{% if o.description %}{{ o.description }}{% elif o.key == 'true' %}the statement holds{% else %}it does not{% endif %}"
|
|
"{% elif o.description is none %}{{ o.key }}{% else %}{{ o.description }}{% endif %}"
|
|
)
|
|
user = (
|
|
"{% set user %}record: {{ state }}" + nl
|
|
+ "question: {{ instructions if instructions else '" + self._DEFAULT_INSTRUCTIONS + "' }}" + nl
|
|
+ "options:{% for o in options %}" + nl + "{{ o.label | trim }}) {{ o.key }}: " + description + "{% endfor %}{% endset %}"
|
|
)
|
|
return (
|
|
"<|im_start|>system" + nl + prompt["system"].strip() + "<|im_end|>" + nl
|
|
+ "<|im_start|>user" + nl + "{% for image in images %}{{ image }}{% endfor %}" + user + "{{ user | trim }}<|im_end|>" + nl
|
|
+ "<|im_start|>assistant" + nl + "<think>{{ '\\n\\n' }}</think>{{ '\\n\\n' }}" + prompt["prefill"]
|
|
)
|
|
|
|
def set_gguf_parameters(self):
|
|
super().set_gguf_parameters()
|
|
self.gguf_writer.add_decision_type(gguf.DecisionType.H2O_LIGHTNING)
|
|
# the answer follows "Answer:", so the token of a label has a leading space
|
|
cfg = self._serve_config
|
|
self.gguf_writer.add_decision_labels([" " + label for label in cfg["labels"][:cfg.get("max_options", len(cfg["labels"]))]])
|
|
for name in ("choice", "noul", "score"):
|
|
self.gguf_writer.add_decision_temperature(name, self._serve_config["temperature"][name])
|
|
# a request with images has this temperature for all question types
|
|
self.gguf_writer.add_decision_temperature("image", self._serve_config["image"]["temperature"])
|
|
|
|
|
|
@ModelBase.register("H2OLightningModel")
|
|
class H2OLightningVisionModel(Qwen3VLVisionModel):
|
|
pass
|