mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-10-28 09:45:55 +08:00
Xuan-Son Nguyen
097cf33d94
* models: support pplx-decider * vllm d2 * update * clarify naming rule * stricter validation * fix ci * rm redundant filter_tensors
108 lines
4.9 KiB
Python
108 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
from pathlib import Path
|
|
from typing import Any, Callable, Iterator, TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from torch import Tensor
|
|
|
|
from .base import ModelBase, gguf, logger
|
|
from .qwen import Qwen3Model
|
|
|
|
|
|
def _is_vllm_decision2_checkpoint(dir_model: Path) -> bool:
|
|
if not (dir_model / "config.json").is_file():
|
|
return False
|
|
with open(dir_model / "config.json", encoding="utf-8") as f:
|
|
return json.load(f).get("decision_format") == "vllm-sr-decision"
|
|
|
|
|
|
@ModelBase.register_hparams_loader(_is_vllm_decision2_checkpoint)
|
|
def _load_vllm_decision2_hparams(dir_model: Path) -> dict[str, Any]:
|
|
logger.info("gguf: detected vllm-sr Decision 2.0 checkpoint")
|
|
with open(dir_model / "config.json", encoding="utf-8") as f:
|
|
package = json.load(f)
|
|
hparams = ModelBase.load_hparams((dir_model / package["backbone"]["config"]).parent, False, guess=False)
|
|
hparams["architectures"] = ["VllmDecision2Model"]
|
|
with open(dir_model / package["model_config"], encoding="utf-8") as f:
|
|
hparams["decision"] = json.load(f)
|
|
with open(dir_model / "score_bias.json", encoding="utf-8") as f:
|
|
hparams["decision"]["score_bias"] = json.load(f)["offsets"]
|
|
return hparams
|
|
|
|
|
|
@ModelBase.register("VllmDecision2Model")
|
|
@ModelBase.example("vllm-sr/Decision-2.0-Kai-0.6B")
|
|
class VllmDecision2Model(Qwen3Model):
|
|
model_arch = gguf.MODEL_ARCH.VLLM_DECISION2
|
|
|
|
# torch.nn.LayerNorm default, used by the head
|
|
_HEAD_NORM_EPS = 1e-5
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
head = self.hparams["decision"]
|
|
if head["architecture"] != "qwen3-text-endpoints-global-query-shared-bilinear-mlp" or head["head_variant"] != "shared":
|
|
raise ValueError(f"unsupported decision head: {head['architecture']}")
|
|
|
|
def index_tensors(self, *args, **kwargs) -> dict[str, Callable[[], Tensor]]:
|
|
# the backbone weights are in their own directory
|
|
dir_model, self.dir_model = self.dir_model, self.dir_model / "backbone"
|
|
try:
|
|
return super().index_tensors(*args, **kwargs)
|
|
finally:
|
|
self.dir_model = dir_model
|
|
|
|
def get_tensors(self) -> Iterator[tuple[str, Tensor]]:
|
|
yield from super().get_tensors()
|
|
from safetensors.torch import load_file
|
|
for name, data in load_file(self.dir_model / "decision_head.safetensors").items():
|
|
yield "decision_head." + name, data
|
|
|
|
def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
|
|
# the head is computed in F32
|
|
if new_name.startswith("decision."):
|
|
return gguf.GGMLQuantizationType.F32
|
|
return super().tensor_force_quant(name, new_name, bid, n_dims)
|
|
|
|
def set_vocab(self):
|
|
super().set_vocab()
|
|
self.gguf_writer.add_chat_template([{"name": "systemone", "template": self._systemone_template()}])
|
|
|
|
@classmethod
|
|
def _systemone_template(cls) -> str:
|
|
# prompt follows decision2/_vendor/dev2model/decision_model.py (segments) of the model repo
|
|
# the pieces of the prompt are tokenized one by one, the server gives the text that separates them (sep)
|
|
# and the text that starts the span of a question or of an option (mark_question, mark_option)
|
|
# the head reads the last token of each option and the last token of the question
|
|
# the start of an option is a piece of its own, so that the spans of two options are not joined
|
|
# the keys of JSON objects are given in sorted order
|
|
def text(value: str) -> str:
|
|
return "{{ " + json.dumps(value) + " }}"
|
|
|
|
def render(name: str) -> str:
|
|
return "{{ " + name + " if " + name + " is string else " + name + " | tojson(separators=[',', ':']) }}"
|
|
|
|
option = (
|
|
"{% set d = o.description %}"
|
|
"{% if q.type == 'noul' and d is none %}{% set d = 'Yes' if o.key == 'true' else 'No' %}{% endif %}"
|
|
"{{ {'description': d, 'key': o.key} | tojson(separators=[',', ':']) }}"
|
|
)
|
|
return (
|
|
"{% for q in questions %}"
|
|
+ text("Context:\n") + render("state")
|
|
+ text("\n\nTask type: ") + "{{ q.type }}" + text("\nQuestion:\n") + render("q.instructions") + text("\nOptions:")
|
|
+ "{% for o in q.options %}{{ sep }}" + text("\n<option>\n") + "{{ sep }}{{ mark_option }}" + option + text("\n</option>") + "{% endfor %}"
|
|
+ "{{ sep }}{{ mark_question }}" + text("\n\nSelect the single option best supported by the context and instructions.\nDecision:")
|
|
+ "{% endfor %}"
|
|
)
|
|
|
|
def set_gguf_parameters(self):
|
|
super().set_gguf_parameters()
|
|
self.gguf_writer.add_decision_type(gguf.DecisionType.VLLM_DECISION2)
|
|
self.gguf_writer.add_layer_norm_eps(self._HEAD_NORM_EPS)
|
|
for offsets in self.hparams["decision"]["score_bias"].values():
|
|
self.gguf_writer.add_decision_score_bias(offsets)
|