341 lines
13 KiB
Python
341 lines
13 KiB
Python
"""Output bundle helpers shared by the GGUF conversion scripts."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import shutil
|
||
from pathlib import Path
|
||
|
||
log = logging.getLogger("nvfp4gguf.bundle")
|
||
|
||
CONTEXT_HPARAMS = (
|
||
"max_position_embeddings",
|
||
"n_ctx",
|
||
"n_positions",
|
||
"max_length",
|
||
"max_sequence_length",
|
||
"model_max_length",
|
||
)
|
||
|
||
|
||
def is_model_context_key(name: str) -> bool:
|
||
"""The model's max context, not RoPE's original context length."""
|
||
return name.endswith(".context_length") and ".rope." not in name
|
||
|
||
|
||
def require_context_length(length: int | None) -> int | None:
|
||
if length is None:
|
||
return None
|
||
if length <= 0:
|
||
raise SystemExit("--context-length must be a positive token count")
|
||
return int(length)
|
||
|
||
|
||
def scalar_context_length(value) -> int | None:
|
||
if isinstance(value, bool) or isinstance(value, (list, dict)):
|
||
return None
|
||
try:
|
||
length = int(value)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
if length <= 0:
|
||
return None
|
||
return length
|
||
|
||
|
||
def _hf_context_key(hparams: dict) -> tuple[str | None, int | None]:
|
||
"""Key llama.cpp's find_hparam will see, and its scalar length if it has one."""
|
||
for key in CONTEXT_HPARAMS:
|
||
if key in hparams:
|
||
return key, scalar_context_length(hparams[key])
|
||
for nested_name in ("text_config", "llm_config", "language_config"):
|
||
nested = hparams.get(nested_name)
|
||
if not isinstance(nested, dict):
|
||
continue
|
||
for key in CONTEXT_HPARAMS:
|
||
if key in nested:
|
||
return key, scalar_context_length(nested[key])
|
||
return None, None
|
||
|
||
|
||
def apply_hf_context_length(hparams: dict, length: int | None) -> int | None:
|
||
"""Return the context llama.cpp will write, after an optional override.
|
||
|
||
llama.cpp reads the first usual context key on the top-level config. A
|
||
list or a missing key is not a usable length. The override is written
|
||
where that lookup will find it.
|
||
"""
|
||
requested = require_context_length(length)
|
||
found_key, found_value = _hf_context_key(hparams)
|
||
if requested is not None:
|
||
hparams[found_key or "max_position_embeddings"] = requested
|
||
log.info("context length %d", requested)
|
||
return requested
|
||
if found_value is None:
|
||
return None
|
||
if found_key not in hparams:
|
||
hparams[found_key] = found_value
|
||
log.info("context length %d from the nested checkpoint config", found_value)
|
||
else:
|
||
log.info("context length %d from the checkpoint", found_value)
|
||
return found_value
|
||
|
||
OUTTYPES = ("auto", "f32", "f16", "bf16", "nvfp4")
|
||
|
||
|
||
def outtype(value: str) -> str:
|
||
return value.lower()
|
||
|
||
|
||
def quote_modelfile_value(text: str) -> str:
|
||
if text == "" or any(char.isspace() for char in text) or '"' in text:
|
||
if '"""' in text:
|
||
raise ValueError("Modelfile value contains a triple quote, which a Modelfile cannot hold")
|
||
if "\n" in text or '"' in text:
|
||
return '"""' + text + '"""'
|
||
return '"' + text + '"'
|
||
return text
|
||
|
||
|
||
def mmproj_destinations(gguf_path: Path, sources: list[Path] | None) -> list[tuple[Path, Path]]:
|
||
"""Pair each projector with the path it will have in the output directory."""
|
||
dest_dir = gguf_path.parent
|
||
planned: list[tuple[Path, Path]] = []
|
||
seen: set[str] = set()
|
||
for source in sources or []:
|
||
source = Path(source)
|
||
if not source.is_file():
|
||
raise SystemExit(f"mmproj file does not exist: {source}")
|
||
if source.name == gguf_path.name or source.name in seen:
|
||
raise SystemExit(
|
||
f"mmproj {source} would be copied as {source.name}, which collides with another file in {dest_dir}."
|
||
)
|
||
seen.add(source.name)
|
||
planned.append((source, dest_dir / source.name))
|
||
return planned
|
||
|
||
|
||
def place_mmprojs(gguf_path: Path, sources: list[Path] | None, overwrite: bool) -> list[Path]:
|
||
"""Copy projector GGUFs next to the language model, keeping their file names."""
|
||
copied: list[Path] = []
|
||
for source, dest in mmproj_destinations(gguf_path, sources):
|
||
try:
|
||
same = dest.exists() and dest.resolve() == source.resolve()
|
||
except OSError:
|
||
same = False
|
||
if same:
|
||
copied.append(dest)
|
||
continue
|
||
if dest.exists() and not overwrite:
|
||
raise SystemExit(f"{dest} already exists. Pass --overwrite to replace it.")
|
||
shutil.copy2(source, dest)
|
||
log.info("Copied vision projector to %s", dest)
|
||
copied.append(dest)
|
||
return copied
|
||
|
||
|
||
def refuse_existing_bundle(gguf_path: Path, overwrite: bool, mmprojs: list[Path] | None, *, skip_gguf: bool = False) -> None:
|
||
"""Stop before a long run when the GGUF, Modelfile, or a copied projector is already there."""
|
||
if not skip_gguf and gguf_path.exists() and not overwrite:
|
||
raise SystemExit(f"{gguf_path} already exists. Pass --overwrite to replace it.")
|
||
modelfile = gguf_path.parent / "Modelfile"
|
||
if modelfile.exists() and not overwrite:
|
||
raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
|
||
for source, dest in mmproj_destinations(gguf_path, mmprojs):
|
||
try:
|
||
same = dest.exists() and dest.resolve() == source.resolve()
|
||
except OSError:
|
||
same = False
|
||
if dest.exists() and not same and not overwrite:
|
||
raise SystemExit(f"{dest} already exists. Pass --overwrite to replace it.")
|
||
|
||
|
||
def _stored_value(kv: dict, name: str):
|
||
item = kv.get(name)
|
||
if item is None:
|
||
return None
|
||
return getattr(item, "value", item)
|
||
|
||
|
||
def _nonempty(value) -> bool:
|
||
if value is None:
|
||
return False
|
||
if isinstance(value, str):
|
||
return bool(value.strip())
|
||
if isinstance(value, (list, tuple)):
|
||
return len(value) > 0
|
||
return True
|
||
|
||
|
||
def warn_missing_gguf_metadata(writer) -> None:
|
||
"""Warn when the GGUF metadata cannot describe a usable chat model.
|
||
|
||
Context length, the chat template, and the tokenizer are the fields that
|
||
still produce a file but change how it runs. Block count and embedding
|
||
length are what llama.cpp needs to build the graph.
|
||
"""
|
||
kv = writer.kv_data[0] if writer.kv_data else {}
|
||
arch = writer.arch or "unknown"
|
||
missing: list[str] = []
|
||
|
||
context_keys = [name for name in kv if is_model_context_key(name)]
|
||
if not context_keys:
|
||
missing.append(f"context length ({arch}.context_length); pass --context-length to set it")
|
||
elif all(scalar_context_length(_stored_value(kv, name)) is None for name in context_keys):
|
||
missing.append("a single context length; pass --context-length to replace the stored value")
|
||
|
||
template_names = [
|
||
name for name in kv
|
||
if name == "tokenizer.chat_template" or name == "tokenizer.chat_templates" or name.startswith("tokenizer.chat_template.")
|
||
]
|
||
if not any(_nonempty(_stored_value(kv, name)) for name in template_names):
|
||
missing.append("chat template (tokenizer.chat_template)")
|
||
|
||
if not _nonempty(_stored_value(kv, "tokenizer.ggml.model")):
|
||
missing.append("tokenizer model (tokenizer.ggml.model)")
|
||
if not _nonempty(_stored_value(kv, "tokenizer.ggml.tokens")):
|
||
missing.append("tokenizer vocabulary (tokenizer.ggml.tokens)")
|
||
if "tokenizer.ggml.eos_token_id" not in kv:
|
||
missing.append("end-of-sequence token (tokenizer.ggml.eos_token_id)")
|
||
if scalar_context_length(_stored_value(kv, f"{arch}.block_count")) is None:
|
||
missing.append(f"block count ({arch}.block_count)")
|
||
if scalar_context_length(_stored_value(kv, f"{arch}.embedding_length")) is None:
|
||
missing.append(f"embedding length ({arch}.embedding_length)")
|
||
|
||
if missing:
|
||
log.warning("GGUF is missing %s.", "; ".join(missing))
|
||
|
||
|
||
# llama.cpp enables its DeepSeek-R1 tool parser only when the template source
|
||
# contains this exact expression. The sglang DeepSeek-R1 template renders the
|
||
# same calls but does not contain the expression, so the calls stay in the
|
||
# visible reply.
|
||
_DEEPSEEK_R1_TOOL_MARKER = (
|
||
"{{'<|Assistant|><|tool▁calls▁begin|><|tool▁call▁begin|>' + tool['type'] + '<|tool▁sep|>'"
|
||
)
|
||
# A brace was left between the two branches of that template. Jinja prints it
|
||
# after every tool call that has no text.
|
||
_DEEPSEEK_R1_STRAY_BRACE = "'<|tool▁call▁end|>'}}\n }\n {%- else %}"
|
||
_DEEPSEEK_R1_STRAY_FIXED = "'<|tool▁call▁end|>'}}\n {%- else %}"
|
||
|
||
|
||
def normalize_chat_template(value):
|
||
"""Make a DeepSeek-R1 tool template recognizable to llama.cpp's parser."""
|
||
if isinstance(value, str):
|
||
return _normalize_one_chat_template(value)
|
||
if isinstance(value, list):
|
||
changed = []
|
||
for item in value:
|
||
if isinstance(item, str):
|
||
changed.append(_normalize_one_chat_template(item))
|
||
elif isinstance(item, dict) and isinstance(item.get("template"), str):
|
||
item = dict(item)
|
||
item["template"] = _normalize_one_chat_template(item["template"])
|
||
changed.append(item)
|
||
else:
|
||
changed.append(item)
|
||
return changed
|
||
return value
|
||
|
||
|
||
def _normalize_one_chat_template(template: str) -> str:
|
||
if "tool▁calls▁begin" not in template:
|
||
return template
|
||
if _DEEPSEEK_R1_STRAY_BRACE in template:
|
||
template = template.replace(_DEEPSEEK_R1_STRAY_BRACE, _DEEPSEEK_R1_STRAY_FIXED, 1)
|
||
if _DEEPSEEK_R1_TOOL_MARKER not in template:
|
||
template = "{# " + _DEEPSEEK_R1_TOOL_MARKER + " #}\n" + template
|
||
log.info("Adjusted the DeepSeek-R1 tool template so tool calls are parsed instead of printed.")
|
||
return template
|
||
|
||
|
||
def flatten_chat_template(value) -> str:
|
||
"""Join a string or a list of named chat templates into one text blob."""
|
||
if value is None:
|
||
return ""
|
||
if isinstance(value, bytes):
|
||
return value.decode("utf-8", "replace")
|
||
if isinstance(value, str):
|
||
return value
|
||
if isinstance(value, list):
|
||
parts: list[str] = []
|
||
for item in value:
|
||
if isinstance(item, str):
|
||
parts.append(item)
|
||
elif isinstance(item, dict) and isinstance(item.get("template"), str):
|
||
parts.append(item["template"])
|
||
return "\n".join(parts)
|
||
return ""
|
||
|
||
|
||
# Ollama accepts a tools request only when the chat template both mentions tools
|
||
# and can render a tool result. A bare "tools" string does not.
|
||
_TOOL_RESULT_MARKERS = (
|
||
"tool_response",
|
||
"tool_results",
|
||
"role'] == 'tool'",
|
||
"role'] == \"tool\"",
|
||
"role\"] == 'tool'",
|
||
"role\"] == \"tool\"",
|
||
"message.role == 'tool'",
|
||
"message.role == \"tool\"",
|
||
"ipython",
|
||
)
|
||
|
||
|
||
def chat_template_has_tools(template: str | None) -> bool:
|
||
"""True when the template can render a tool call and the tool's result."""
|
||
if not template:
|
||
return False
|
||
if "tools" not in template and "tool_call" not in template:
|
||
return False
|
||
if "tool_calls" not in template and "assistant_tool_call" not in template:
|
||
return False
|
||
return any(marker in template for marker in _TOOL_RESULT_MARKERS)
|
||
|
||
|
||
def tool_modelfile_lines(template: str | None) -> list[str]:
|
||
"""Declare tools when the stored chat template can round-trip them.
|
||
|
||
A template that cannot render a tool result leaves the capability off.
|
||
Ollama then rejects a request that includes tools instead of dropping the
|
||
call on the floor.
|
||
"""
|
||
if chat_template_has_tools(template):
|
||
return [
|
||
"CAPABILITY tools",
|
||
"# Tool calls use the chat template stored in the GGUF.",
|
||
]
|
||
return ["# This chat template has no tool round-trip, so tools stay off."]
|
||
|
||
|
||
def write_simple_modelfile(
|
||
gguf_path: Path,
|
||
overwrite: bool,
|
||
mmprojs: list[Path] | None,
|
||
chat_template: str | None = None,
|
||
) -> Path:
|
||
"""Write a Modelfile that points at a finished GGUF and any projector copies.
|
||
|
||
Context is not set here. A PARAMETER num_ctx would override OLLAMA_CONTEXT_LENGTH.
|
||
"""
|
||
modelfile = gguf_path.parent / "Modelfile"
|
||
if modelfile.exists() and not overwrite:
|
||
raise SystemExit(f"{modelfile} already exists. Pass --overwrite to replace it.")
|
||
projectors = place_mmprojs(gguf_path, mmprojs, overwrite)
|
||
lines = [f"FROM ./{gguf_path.name}"]
|
||
for projector in projectors:
|
||
lines.append(f"FROM ./{projector.name}")
|
||
lines.append("")
|
||
lines.append("# The chat template is the one stored in the GGUF.")
|
||
lines.extend(tool_modelfile_lines(chat_template))
|
||
if chat_template_has_tools(chat_template):
|
||
log.info("Chat template supports tools.")
|
||
elif chat_template:
|
||
log.info("Chat template has no tool round-trip. Tool requests stay rejected.")
|
||
lines.append("# Context length comes from OLLAMA_CONTEXT_LENGTH.")
|
||
lines.append("")
|
||
modelfile.write_text("\n".join(lines), encoding="utf-8")
|
||
log.info("Wrote %s", modelfile)
|
||
return modelfile
|