summaryrefslogtreecommitdiff
path: root/gguf-py
diff options
context:
space:
mode:
Diffstat (limited to 'gguf-py')
-rw-r--r--gguf-py/README.md19
-rw-r--r--gguf-py/gguf/__init__.py2
-rw-r--r--gguf-py/gguf/constants.py569
-rw-r--r--gguf-py/gguf/gguf_reader.py35
-rw-r--r--gguf-py/gguf/gguf_writer.py445
-rw-r--r--gguf-py/gguf/lazy.py71
-rw-r--r--gguf-py/gguf/metadata.py503
-rw-r--r--gguf-py/gguf/quants.py2
-rw-r--r--gguf-py/gguf/tensor_mapping.py192
-rw-r--r--gguf-py/gguf/utility.py69
-rw-r--r--gguf-py/pyproject.toml3
-rw-r--r--gguf-py/scripts/__init__.py17
-rwxr-xr-xgguf-py/scripts/gguf_convert_endian.py (renamed from gguf-py/scripts/gguf-convert-endian.py)0
-rwxr-xr-xgguf-py/scripts/gguf_dump.py (renamed from gguf-py/scripts/gguf-dump.py)80
-rwxr-xr-xgguf-py/scripts/gguf_hash.py102
-rwxr-xr-xgguf-py/scripts/gguf_new_metadata.py (renamed from gguf-py/scripts/gguf-new-metadata.py)2
-rwxr-xr-xgguf-py/scripts/gguf_set_metadata.py (renamed from gguf-py/scripts/gguf-set-metadata.py)0
-rw-r--r--gguf-py/tests/__init__.py1
-rw-r--r--gguf-py/tests/test_gguf.py7
-rwxr-xr-xgguf-py/tests/test_metadata.py203
20 files changed, 1941 insertions, 381 deletions
diff --git a/gguf-py/README.md b/gguf-py/README.md
index a04c2275..24af96a1 100644
--- a/gguf-py/README.md
+++ b/gguf-py/README.md
@@ -3,7 +3,7 @@
This is a Python package for writing binary files in the [GGUF](https://github.com/ggerganov/ggml/pull/302)
(GGML Universal File) format.
-See [convert-llama-hf-to-gguf.py](https://github.com/ggerganov/llama.cpp/blob/master/convert-hf-to-gguf.py)
+See [convert_hf_to_gguf.py](https://github.com/ggerganov/llama.cpp/blob/master/convert_hf_to_gguf.py)
as an example for its usage.
## Installation
@@ -15,13 +15,13 @@ pip install gguf
[examples/writer.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/examples/writer.py) — Generates `example.gguf` in the current directory to demonstrate generating a GGUF file. Note that this file cannot be used as a model.
-[scripts/gguf-dump.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf-dump.py) — Dumps a GGUF file's metadata to the console.
+[scripts/gguf_dump.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf_dump.py) — Dumps a GGUF file's metadata to the console.
-[scripts/gguf-set-metadata.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf-set-metadata.py) — Allows changing simple metadata values in a GGUF file by key.
+[scripts/gguf_set_metadata.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf_set_metadata.py) — Allows changing simple metadata values in a GGUF file by key.
-[scripts/gguf-convert-endian.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf-convert-endian.py) — Allows converting the endianness of GGUF files.
+[scripts/gguf_convert_endian.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf_convert_endian.py) — Allows converting the endianness of GGUF files.
-[scripts/gguf-new-metadata.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf-new-metadata.py) — Copies a GGUF file with added/modified/removed metadata values.
+[scripts/gguf_new_metadata.py](https://github.com/ggerganov/llama.cpp/blob/master/gguf-py/scripts/gguf_new_metadata.py) — Copies a GGUF file with added/modified/removed metadata values.
## Development
Maintainers who participate in development of this package are advised to install it in editable mode:
@@ -78,6 +78,13 @@ python -m build
python -m twine upload dist/*
```
+## Run Unit Tests
+
+From root of this repository you can run this command to run all the unit tests
+
+```bash
+python -m unittest discover ./gguf-py -v
+```
+
## TODO
-- [ ] Add tests
- [ ] Include conversion scripts as command line entry points in this package.
diff --git a/gguf-py/gguf/__init__.py b/gguf-py/gguf/__init__.py
index ea5146b1..243defc4 100644
--- a/gguf-py/gguf/__init__.py
+++ b/gguf-py/gguf/__init__.py
@@ -5,3 +5,5 @@ from .gguf_writer import *
from .quants import *
from .tensor_mapping import *
from .vocab import *
+from .utility import *
+from .metadata import *
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 4cc3e35f..e343c2ef 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -19,18 +19,60 @@ GGML_QUANT_VERSION = 2 # GGML_QNT_VERSION from ggml.h
class Keys:
class General:
- ARCHITECTURE = "general.architecture"
- QUANTIZATION_VERSION = "general.quantization_version"
- ALIGNMENT = "general.alignment"
- NAME = "general.name"
- AUTHOR = "general.author"
- VERSION = "general.version"
- URL = "general.url"
- DESCRIPTION = "general.description"
- LICENSE = "general.license"
- SOURCE_URL = "general.source.url"
- SOURCE_HF_REPO = "general.source.huggingface.repository"
- FILE_TYPE = "general.file_type"
+ TYPE = "general.type"
+ ARCHITECTURE = "general.architecture"
+ QUANTIZATION_VERSION = "general.quantization_version"
+ ALIGNMENT = "general.alignment"
+ FILE_TYPE = "general.file_type"
+
+ # Authorship Metadata
+ NAME = "general.name"
+ AUTHOR = "general.author"
+ VERSION = "general.version"
+ ORGANIZATION = "general.organization"
+
+ FINETUNE = "general.finetune"
+ BASENAME = "general.basename"
+
+ DESCRIPTION = "general.description"
+ QUANTIZED_BY = "general.quantized_by"
+
+ SIZE_LABEL = "general.size_label"
+
+ # Licensing details
+ LICENSE = "general.license"
+ LICENSE_NAME = "general.license.name"
+ LICENSE_LINK = "general.license.link"
+
+ # Typically represents the converted GGUF repo (Unless native)
+ URL = "general.url" # Model Website/Paper
+ DOI = "general.doi"
+ UUID = "general.uuid"
+ REPO_URL = "general.repo_url" # Model Source Repository (git/svn/etc...)
+
+ # Model Source during conversion
+ SOURCE_URL = "general.source.url" # Model Website/Paper
+ SOURCE_DOI = "general.source.doi"
+ SOURCE_UUID = "general.source.uuid"
+ SOURCE_REPO_URL = "general.source.repo_url" # Model Source Repository (git/svn/etc...)
+
+ # Base Model Source. There can be more than one source if it's a merged
+ # model like with 'Mistral-7B-Merge-14-v0.1'. This will assist in
+ # tracing linage of models as it is finetuned or merged over time.
+ BASE_MODEL_COUNT = "general.base_model.count"
+ BASE_MODEL_NAME = "general.base_model.{id}.name"
+ BASE_MODEL_AUTHOR = "general.base_model.{id}.author"
+ BASE_MODEL_VERSION = "general.base_model.{id}.version"
+ BASE_MODEL_ORGANIZATION = "general.base_model.{id}.organization"
+ BASE_MODEL_URL = "general.base_model.{id}.url" # Model Website/Paper
+ BASE_MODEL_DOI = "general.base_model.{id}.doi"
+ BASE_MODEL_UUID = "general.base_model.{id}.uuid"
+ BASE_MODEL_REPO_URL = "general.base_model.{id}.repo_url" # Model Source Repository (git/svn/etc...)
+
+ # Array based KV stores
+ TAGS = "general.tags"
+ LANGUAGES = "general.languages"
+ DATASETS = "general.datasets"
class LLM:
VOCAB_SIZE = "{arch}.vocab_size"
@@ -49,6 +91,9 @@ class Keys:
EXPERT_WEIGHTS_SCALE = "{arch}.expert_weights_scale"
POOLING_TYPE = "{arch}.pooling_type"
LOGIT_SCALE = "{arch}.logit_scale"
+ DECODER_START_TOKEN_ID = "{arch}.decoder_start_token_id"
+ ATTN_LOGIT_SOFTCAPPING = "{arch}.attn_logit_softcapping"
+ FINAL_LOGIT_SOFTCAPPING = "{arch}.final_logit_softcapping"
class Attention:
HEAD_COUNT = "{arch}.attention.head_count"
@@ -62,6 +107,8 @@ class Keys:
CAUSAL = "{arch}.attention.causal"
Q_LORA_RANK = "{arch}.attention.q_lora_rank"
KV_LORA_RANK = "{arch}.attention.kv_lora_rank"
+ REL_BUCKETS_COUNT = "{arch}.attention.relative_buckets_count"
+ SLIDING_WINDOW = "{arch}.attention.sliding_window"
class Rope:
DIMENSION_COUNT = "{arch}.rope.dimension_count"
@@ -73,6 +120,11 @@ class Keys:
SCALING_FINETUNED = "{arch}.rope.scaling.finetuned"
SCALING_YARN_LOG_MUL = "{arch}.rope.scaling.yarn_log_multiplier"
+ class Split:
+ LLM_KV_SPLIT_NO = "split.no"
+ LLM_KV_SPLIT_COUNT = "split.count"
+ LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"
+
class SSM:
CONV_KERNEL = "{arch}.ssm.conv_kernel"
INNER_SIZE = "{arch}.ssm.inner_size"
@@ -80,129 +132,175 @@ class Keys:
TIME_STEP_RANK = "{arch}.ssm.time_step_rank"
class Tokenizer:
- MODEL = "tokenizer.ggml.model"
- PRE = "tokenizer.ggml.pre"
- LIST = "tokenizer.ggml.tokens"
- TOKEN_TYPE = "tokenizer.ggml.token_type"
- TOKEN_TYPE_COUNT = "tokenizer.ggml.token_type_count" # for BERT-style token types
- SCORES = "tokenizer.ggml.scores"
- MERGES = "tokenizer.ggml.merges"
- BOS_ID = "tokenizer.ggml.bos_token_id"
- EOS_ID = "tokenizer.ggml.eos_token_id"
- UNK_ID = "tokenizer.ggml.unknown_token_id"
- SEP_ID = "tokenizer.ggml.seperator_token_id"
- PAD_ID = "tokenizer.ggml.padding_token_id"
- CLS_ID = "tokenizer.ggml.cls_token_id"
- MASK_ID = "tokenizer.ggml.mask_token_id"
- ADD_BOS = "tokenizer.ggml.add_bos_token"
- ADD_EOS = "tokenizer.ggml.add_eos_token"
- ADD_PREFIX = "tokenizer.ggml.add_space_prefix"
- HF_JSON = "tokenizer.huggingface.json"
- RWKV = "tokenizer.rwkv.world"
- CHAT_TEMPLATE = "tokenizer.chat_template"
- CHAT_TEMPLATE_N = "tokenizer.chat_template.{name}"
- CHAT_TEMPLATES = "tokenizer.chat_templates"
+ MODEL = "tokenizer.ggml.model"
+ PRE = "tokenizer.ggml.pre"
+ LIST = "tokenizer.ggml.tokens"
+ TOKEN_TYPE = "tokenizer.ggml.token_type"
+ TOKEN_TYPE_COUNT = "tokenizer.ggml.token_type_count" # for BERT-style token types
+ SCORES = "tokenizer.ggml.scores"
+ MERGES = "tokenizer.ggml.merges"
+ BOS_ID = "tokenizer.ggml.bos_token_id"
+ EOS_ID = "tokenizer.ggml.eos_token_id"
+ UNK_ID = "tokenizer.ggml.unknown_token_id"
+ SEP_ID = "tokenizer.ggml.seperator_token_id"
+ PAD_ID = "tokenizer.ggml.padding_token_id"
+ CLS_ID = "tokenizer.ggml.cls_token_id"
+ MASK_ID = "tokenizer.ggml.mask_token_id"
+ ADD_BOS = "tokenizer.ggml.add_bos_token"
+ ADD_EOS = "tokenizer.ggml.add_eos_token"
+ ADD_PREFIX = "tokenizer.ggml.add_space_prefix"
+ REMOVE_EXTRA_WS = "tokenizer.ggml.remove_extra_whitespaces"
+ PRECOMPILED_CHARSMAP = "tokenizer.ggml.precompiled_charsmap"
+ HF_JSON = "tokenizer.huggingface.json"
+ RWKV = "tokenizer.rwkv.world"
+ CHAT_TEMPLATE = "tokenizer.chat_template"
+ CHAT_TEMPLATE_N = "tokenizer.chat_template.{name}"
+ CHAT_TEMPLATES = "tokenizer.chat_templates"
# FIM/Infill special tokens constants
- PREFIX_ID = "tokenizer.ggml.prefix_token_id"
- SUFFIX_ID = "tokenizer.ggml.suffix_token_id"
- MIDDLE_ID = "tokenizer.ggml.middle_token_id"
- EOT_ID = "tokenizer.ggml.eot_token_id"
+ PREFIX_ID = "tokenizer.ggml.prefix_token_id"
+ SUFFIX_ID = "tokenizer.ggml.suffix_token_id"
+ MIDDLE_ID = "tokenizer.ggml.middle_token_id"
+ EOT_ID = "tokenizer.ggml.eot_token_id"
+ class Adapter:
+ TYPE = "adapter.type"
+ LORA_ALPHA = "adapter.lora.alpha"
#
# recommended mapping of model tensor names for storage in gguf
#
+class GGUFType:
+ MODEL = "model"
+ ADAPTER = "adapter"
+
+
class MODEL_ARCH(IntEnum):
- LLAMA = auto()
- FALCON = auto()
- BAICHUAN = auto()
- GROK = auto()
- GPT2 = auto()
- GPTJ = auto()
- GPTNEOX = auto()
- MPT = auto()
- STARCODER = auto()
- REFACT = auto()
- BERT = auto()
- NOMIC_BERT = auto()
+ LLAMA = auto()
+ FALCON = auto()
+ BAICHUAN = auto()
+ GROK = auto()
+ GPT2 = auto()
+ GPTJ = auto()
+ GPTNEOX = auto()
+ MPT = auto()
+ STARCODER = auto()
+ REFACT = auto()
+ BERT = auto()
+ NOMIC_BERT = auto()
JINA_BERT_V2 = auto()
- BLOOM = auto()
- STABLELM = auto()
- QWEN = auto()
- QWEN2 = auto()
- QWEN2MOE = auto()
- PHI2 = auto()
- PHI3 = auto()
- PLAMO = auto()
- CODESHELL = auto()
- ORION = auto()
- INTERNLM2 = auto()
- MINICPM = auto()
- GEMMA = auto()
- STARCODER2 = auto()
- MAMBA = auto()
- XVERSE = auto()
- COMMAND_R = auto()
- DBRX = auto()
- OLMO = auto()
- ARCTIC = auto()
- DEEPSEEK2 = auto()
- BITNET = auto()
+ BLOOM = auto()
+ STABLELM = auto()
+ QWEN = auto()
+ QWEN2 = auto()
+ QWEN2MOE = auto()
+ PHI2 = auto()
+ PHI3 = auto()
+ PLAMO = auto()
+ CODESHELL = auto()
+ ORION = auto()
+ INTERNLM2 = auto()
+ MINICPM = auto()
+ GEMMA = auto()
+ GEMMA2 = auto()
+ STARCODER2 = auto()
+ MAMBA = auto()
+ XVERSE = auto()
+ COMMAND_R = auto()
+ DBRX = auto()
+ OLMO = auto()
+ OPENELM = auto()
+ ARCTIC = auto()
+ DEEPSEEK2 = auto()
+ CHATGLM = auto()
+ BITNET = auto()
+ T5 = auto()
+ JAIS = auto()
class MODEL_TENSOR(IntEnum):
- TOKEN_EMBD = auto()
- TOKEN_EMBD_NORM = auto()
- TOKEN_TYPES = auto()
- POS_EMBD = auto()
- OUTPUT = auto()
- OUTPUT_NORM = auto()
- ROPE_FREQS = auto()
- ROPE_FACTORS_LONG = auto()
- ROPE_FACTORS_SHORT = auto()
- ATTN_Q = auto()
- ATTN_K = auto()
- ATTN_V = auto()
- ATTN_QKV = auto()
- ATTN_OUT = auto()
- ATTN_NORM = auto()
- ATTN_NORM_2 = auto()
- ATTN_OUT_NORM = auto()
- ATTN_ROT_EMBD = auto()
- FFN_GATE_INP = auto()
- FFN_GATE_INP_SHEXP = auto()
- FFN_NORM = auto()
- FFN_GATE = auto()
- FFN_DOWN = auto()
- FFN_UP = auto()
- FFN_ACT = auto()
- FFN_NORM_EXP = auto()
- FFN_GATE_EXP = auto()
- FFN_DOWN_EXP = auto()
- FFN_UP_EXP = auto()
- FFN_GATE_SHEXP = auto()
- FFN_DOWN_SHEXP = auto()
- FFN_UP_SHEXP = auto()
- ATTN_Q_NORM = auto()
- ATTN_K_NORM = auto()
- LAYER_OUT_NORM = auto()
- SSM_IN = auto()
- SSM_CONV1D = auto()
- SSM_X = auto()
- SSM_DT = auto()
- SSM_A = auto()
- SSM_D = auto()
- SSM_OUT = auto()
- ATTN_Q_A = auto()
- ATTN_Q_B = auto()
- ATTN_KV_A_MQA = auto()
- ATTN_KV_B = auto()
- ATTN_Q_A_NORM = auto()
- ATTN_KV_A_NORM = auto()
- FFN_SUB_NORM = auto()
- ATTN_SUB_NORM = auto()
+ TOKEN_EMBD = auto()
+ TOKEN_EMBD_NORM = auto()
+ TOKEN_TYPES = auto()
+ POS_EMBD = auto()
+ OUTPUT = auto()
+ OUTPUT_NORM = auto()
+ ROPE_FREQS = auto()
+ ROPE_FACTORS_LONG = auto()
+ ROPE_FACTORS_SHORT = auto()
+ ATTN_Q = auto()
+ ATTN_K = auto()
+ ATTN_V = auto()
+ ATTN_QKV = auto()
+ ATTN_OUT = auto()
+ ATTN_NORM = auto()
+ ATTN_NORM_2 = auto()
+ ATTN_OUT_NORM = auto()
+ ATTN_POST_NORM = auto()
+ ATTN_ROT_EMBD = auto()
+ FFN_GATE_INP = auto()
+ FFN_GATE_INP_SHEXP = auto()
+ FFN_NORM = auto()
+ FFN_PRE_NORM = auto()
+ FFN_POST_NORM = auto()
+ FFN_GATE = auto()
+ FFN_DOWN = auto()
+ FFN_UP = auto()
+ FFN_ACT = auto()
+ FFN_NORM_EXP = auto()
+ FFN_GATE_EXP = auto()
+ FFN_DOWN_EXP = auto()
+ FFN_UP_EXP = auto()
+ FFN_GATE_SHEXP = auto()
+ FFN_DOWN_SHEXP = auto()
+ FFN_UP_SHEXP = auto()
+ ATTN_Q_NORM = auto()
+ ATTN_K_NORM = auto()
+ LAYER_OUT_NORM = auto()
+ SSM_IN = auto()
+ SSM_CONV1D = auto()
+ SSM_X = auto()
+ SSM_DT = auto()
+ SSM_A = auto()
+ SSM_D = auto()
+ SSM_OUT = auto()
+ ATTN_Q_A = auto()
+ ATTN_Q_B = auto()
+ ATTN_KV_A_MQA = auto()
+ ATTN_KV_B = auto()
+ ATTN_Q_A_NORM = auto()
+ ATTN_KV_A_NORM = auto()
+ FFN_SUB_NORM = auto()
+ ATTN_SUB_NORM = auto()
+ DEC_ATTN_NORM = auto()
+ DEC_ATTN_Q = auto()
+ DEC_ATTN_K = auto()
+ DEC_ATTN_V = auto()
+ DEC_ATTN_OUT = auto()
+ DEC_ATTN_REL_B = auto()
+ DEC_CROSS_ATTN_NORM = auto()
+ DEC_CROSS_ATTN_Q = auto()
+ DEC_CROSS_ATTN_K = auto()
+ DEC_CROSS_ATTN_V = auto()
+ DEC_CROSS_ATTN_OUT = auto()
+ DEC_CROSS_ATTN_REL_B = auto()
+ DEC_FFN_NORM = auto()
+ DEC_FFN_GATE = auto()
+ DEC_FFN_DOWN = auto()
+ DEC_FFN_UP = auto()
+ DEC_OUTPUT_NORM = auto()
+ ENC_ATTN_NORM = auto()
+ ENC_ATTN_Q = auto()
+ ENC_ATTN_K = auto()
+ ENC_ATTN_V = auto()
+ ENC_ATTN_OUT = auto()
+ ENC_ATTN_REL_B = auto()
+ ENC_FFN_NORM = auto()
+ ENC_FFN_GATE = auto()
+ ENC_FFN_DOWN = auto()
+ ENC_FFN_UP = auto()
+ ENC_OUTPUT_NORM = auto()
MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
@@ -232,68 +330,104 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.INTERNLM2: "internlm2",
MODEL_ARCH.MINICPM: "minicpm",
MODEL_ARCH.GEMMA: "gemma",
+ MODEL_ARCH.GEMMA2: "gemma2",
MODEL_ARCH.STARCODER2: "starcoder2",
MODEL_ARCH.MAMBA: "mamba",
MODEL_ARCH.XVERSE: "xverse",
MODEL_ARCH.COMMAND_R: "command-r",
MODEL_ARCH.DBRX: "dbrx",
MODEL_ARCH.OLMO: "olmo",
+ MODEL_ARCH.OPENELM: "openelm",
MODEL_ARCH.ARCTIC: "arctic",
MODEL_ARCH.DEEPSEEK2: "deepseek2",
+ MODEL_ARCH.CHATGLM: "chatglm",
MODEL_ARCH.BITNET: "bitnet",
+ MODEL_ARCH.T5: "t5",
+ MODEL_ARCH.JAIS: "jais",
}
TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
- MODEL_TENSOR.TOKEN_EMBD: "token_embd",
- MODEL_TENSOR.TOKEN_EMBD_NORM: "token_embd_norm",
- MODEL_TENSOR.TOKEN_TYPES: "token_types",
- MODEL_TENSOR.POS_EMBD: "position_embd",
- MODEL_TENSOR.OUTPUT_NORM: "output_norm",
- MODEL_TENSOR.OUTPUT: "output",
- MODEL_TENSOR.ROPE_FREQS: "rope_freqs",
- MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long",
- MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short",
- MODEL_TENSOR.ATTN_NORM: "blk.{bid}.attn_norm",
- MODEL_TENSOR.ATTN_NORM_2: "blk.{bid}.attn_norm_2",
- MODEL_TENSOR.ATTN_QKV: "blk.{bid}.attn_qkv",
- MODEL_TENSOR.ATTN_Q: "blk.{bid}.attn_q",
- MODEL_TENSOR.ATTN_K: "blk.{bid}.attn_k",
- MODEL_TENSOR.ATTN_V: "blk.{bid}.attn_v",
- MODEL_TENSOR.ATTN_OUT: "blk.{bid}.attn_output",
- MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd",
- MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm",
- MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm",
- MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm",
- MODEL_TENSOR.FFN_GATE_INP: "blk.{bid}.ffn_gate_inp",
- MODEL_TENSOR.FFN_GATE_INP_SHEXP: "blk.{bid}.ffn_gate_inp_shexp",
- MODEL_TENSOR.FFN_NORM: "blk.{bid}.ffn_norm",
- MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate",
- MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down",
- MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up",
- MODEL_TENSOR.FFN_GATE_SHEXP: "blk.{bid}.ffn_gate_shexp",
- MODEL_TENSOR.FFN_DOWN_SHEXP: "blk.{bid}.ffn_down_shexp",
- MODEL_TENSOR.FFN_UP_SHEXP: "blk.{bid}.ffn_up_shexp",
- MODEL_TENSOR.FFN_ACT: "blk.{bid}.ffn",
- MODEL_TENSOR.FFN_NORM_EXP: "blk.{bid}.ffn_norm_exps",
- MODEL_TENSOR.FFN_GATE_EXP: "blk.{bid}.ffn_gate_exps",
- MODEL_TENSOR.FFN_DOWN_EXP: "blk.{bid}.ffn_down_exps",
- MODEL_TENSOR.FFN_UP_EXP: "blk.{bid}.ffn_up_exps",
- MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm",
- MODEL_TENSOR.SSM_IN: "blk.{bid}.ssm_in",
- MODEL_TENSOR.SSM_CONV1D: "blk.{bid}.ssm_conv1d",
- MODEL_TENSOR.SSM_X: "blk.{bid}.ssm_x",
- MODEL_TENSOR.SSM_DT: "blk.{bid}.ssm_dt",
- MODEL_TENSOR.SSM_A: "blk.{bid}.ssm_a",
- MODEL_TENSOR.SSM_D: "blk.{bid}.ssm_d",
- MODEL_TENSOR.SSM_OUT: "blk.{bid}.ssm_out",
- MODEL_TENSOR.ATTN_Q_A: "blk.{bid}.attn_q_a",
- MODEL_TENSOR.ATTN_Q_B: "blk.{bid}.attn_q_b",
- MODEL_TENSOR.ATTN_KV_A_MQA: "blk.{bid}.attn_kv_a_mqa",
- MODEL_TENSOR.ATTN_KV_B: "blk.{bid}.attn_kv_b",
- MODEL_TENSOR.ATTN_Q_A_NORM: "blk.{bid}.attn_q_a_norm",
- MODEL_TENSOR.ATTN_KV_A_NORM: "blk.{bid}.attn_kv_a_norm",
- MODEL_TENSOR.ATTN_SUB_NORM: "blk.{bid}.attn_sub_norm",
- MODEL_TENSOR.FFN_SUB_NORM: "blk.{bid}.ffn_sub_norm",
+ MODEL_TENSOR.TOKEN_EMBD: "token_embd",
+ MODEL_TENSOR.TOKEN_EMBD_NORM: "token_embd_norm",
+ MODEL_TENSOR.TOKEN_TYPES: "token_types",
+ MODEL_TENSOR.POS_EMBD: "position_embd",
+ MODEL_TENSOR.OUTPUT_NORM: "output_norm",
+ MODEL_TENSOR.OUTPUT: "output",
+ MODEL_TENSOR.ROPE_FREQS: "rope_freqs",
+ MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long",
+ MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short",
+ MODEL_TENSOR.ATTN_NORM: "blk.{bid}.attn_norm",
+ MODEL_TENSOR.ATTN_NORM_2: "blk.{bid}.attn_norm_2",
+ MODEL_TENSOR.ATTN_QKV: "blk.{bid}.attn_qkv",
+ MODEL_TENSOR.ATTN_Q: "blk.{bid}.attn_q",
+ MODEL_TENSOR.ATTN_K: "blk.{bid}.attn_k",
+ MODEL_TENSOR.ATTN_V: "blk.{bid}.attn_v",
+ MODEL_TENSOR.ATTN_OUT: "blk.{bid}.attn_output",
+ MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd",
+ MODEL_TENSOR.ATTN_Q_NORM: "blk.{bid}.attn_q_norm",
+ MODEL_TENSOR.ATTN_K_NORM: "blk.{bid}.attn_k_norm",
+ MODEL_TENSOR.ATTN_OUT_NORM: "blk.{bid}.attn_output_norm",
+ MODEL_TENSOR.ATTN_POST_NORM: "blk.{bid}.post_attention_norm",
+ MODEL_TENSOR.FFN_GATE_INP: "blk.{bid}.ffn_gate_inp",
+ MODEL_TENSOR.FFN_GATE_INP_SHEXP: "blk.{bid}.ffn_gate_inp_shexp",
+ MODEL_TENSOR.FFN_NORM: "blk.{bid}.ffn_norm",
+ MODEL_TENSOR.FFN_PRE_NORM: "blk.{bid}.ffn_norm",
+ MODEL_TENSOR.FFN_POST_NORM: "blk.{bid}.post_ffw_norm",
+ MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate",
+ MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down",
+ MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up",
+ MODEL_TENSOR.FFN_GATE_SHEXP: "blk.{bid}.ffn_gate_shexp",
+ MODEL_TENSOR.FFN_DOWN_SHEXP: "blk.{bid}.ffn_down_shexp",
+ MODEL_TENSOR.FFN_UP_SHEXP: "blk.{bid}.ffn_up_shexp",
+ MODEL_TENSOR.FFN_ACT: "blk.{bid}.ffn",
+ MODEL_TENSOR.FFN_NORM_EXP: "blk.{bid}.ffn_norm_exps",
+ MODEL_TENSOR.FFN_GATE_EXP: "blk.{bid}.ffn_gate_exps",
+ MODEL_TENSOR.FFN_DOWN_EXP: "blk.{bid}.ffn_down_exps",
+ MODEL_TENSOR.FFN_UP_EXP: "blk.{bid}.ffn_up_exps",
+ MODEL_TENSOR.LAYER_OUT_NORM: "blk.{bid}.layer_output_norm",
+ MODEL_TENSOR.SSM_IN: "blk.{bid}.ssm_in",
+ MODEL_TENSOR.SSM_CONV1D: "blk.{bid}.ssm_conv1d",
+ MODEL_TENSOR.SSM_X: "blk.{bid}.ssm_x",
+ MODEL_TENSOR.SSM_DT: "blk.{bid}.ssm_dt",
+ MODEL_TENSOR.SSM_A: "blk.{bid}.ssm_a",
+ MODEL_TENSOR.SSM_D: "blk.{bid}.ssm_d",
+ MODEL_TENSOR.SSM_OUT: "blk.{bid}.ssm_out",
+ MODEL_TENSOR.ATTN_Q_A: "blk.{bid}.attn_q_a",
+ MODEL_TENSOR.ATTN_Q_B: "blk.{bid}.attn_q_b",
+ MODEL_TENSOR.ATTN_KV_A_MQA: "blk.{bid}.attn_kv_a_mqa",
+ MODEL_TENSOR.ATTN_KV_B: "blk.{bid}.attn_kv_b",
+ MODEL_TENSOR.ATTN_Q_A_NORM: "blk.{bid}.attn_q_a_norm",
+ MODEL_TENSOR.ATTN_KV_A_NORM: "blk.{bid}.attn_kv_a_norm",
+ MODEL_TENSOR.ATTN_SUB_NORM: "blk.{bid}.attn_sub_norm",
+ MODEL_TENSOR.FFN_SUB_NORM: "blk.{bid}.ffn_sub_norm",
+ MODEL_TENSOR.DEC_ATTN_NORM: "dec.blk.{bid}.attn_norm",
+ MODEL_TENSOR.DEC_ATTN_Q: "dec.blk.{bid}.attn_q",
+ MODEL_TENSOR.DEC_ATTN_K: "dec.blk.{bid}.attn_k",
+ MODEL_TENSOR.DEC_ATTN_V: "dec.blk.{bid}.attn_v",
+ MODEL_TENSOR.DEC_ATTN_OUT: "dec.blk.{bid}.attn_o",
+ MODEL_TENSOR.DEC_ATTN_REL_B: "dec.blk.{bid}.attn_rel_b",
+ MODEL_TENSOR.DEC_CROSS_ATTN_NORM: "dec.blk.{bid}.cross_attn_norm",
+ MODEL_TENSOR.DEC_CROSS_ATTN_Q: "dec.blk.{bid}.cross_attn_q",
+ MODEL_TENSOR.DEC_CROSS_ATTN_K: "dec.blk.{bid}.cross_attn_k",
+ MODEL_TENSOR.DEC_CROSS_ATTN_V: "dec.blk.{bid}.cross_attn_v",
+ MODEL_TENSOR.DEC_CROSS_ATTN_OUT: "dec.blk.{bid}.cross_attn_o",
+ MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: "dec.blk.{bid}.cross_attn_rel_b",
+ MODEL_TENSOR.DEC_FFN_NORM: "dec.blk.{bid}.ffn_norm",
+ MODEL_TENSOR.DEC_FFN_GATE: "dec.blk.{bid}.ffn_gate",
+ MODEL_TENSOR.DEC_FFN_DOWN: "dec.blk.{bid}.ffn_down",
+ MODEL_TENSOR.DEC_FFN_UP: "dec.blk.{bid}.ffn_up",
+ MODEL_TENSOR.DEC_OUTPUT_NORM: "dec.output_norm",
+ MODEL_TENSOR.ENC_ATTN_NORM: "enc.blk.{bid}.attn_norm",
+ MODEL_TENSOR.ENC_ATTN_Q: "enc.blk.{bid}.attn_q",
+ MODEL_TENSOR.ENC_ATTN_K: "enc.blk.{bid}.attn_k",
+ MODEL_TENSOR.ENC_ATTN_V: "enc.blk.{bid}.attn_v",
+ MODEL_TENSOR.ENC_ATTN_OUT: "enc.blk.{bid}.attn_o",
+ MODEL_TENSOR.ENC_ATTN_REL_B: "enc.blk.{bid}.attn_rel_b",
+ MODEL_TENSOR.ENC_FFN_NORM: "enc.blk.{bid}.ffn_norm",
+ MODEL_TENSOR.ENC_FFN_GATE: "enc.blk.{bid}.ffn_gate",
+ MODEL_TENSOR.ENC_FFN_DOWN: "enc.blk.{bid}.ffn_down",
+ MODEL_TENSOR.ENC_FFN_UP: "enc.blk.{bid}.ffn_up",
+ MODEL_TENSOR.ENC_OUTPUT_NORM: "enc.output_norm",
}
MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
@@ -684,6 +818,21 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_UP,
MODEL_TENSOR.FFN_NORM,
],
+ MODEL_ARCH.GEMMA2: [
+ MODEL_TENSOR.TOKEN_EMBD,
+ MODEL_TENSOR.OUTPUT_NORM,
+ MODEL_TENSOR.ATTN_Q,
+ MODEL_TENSOR.ATTN_K,
+ MODEL_TENSOR.ATTN_V,
+ MODEL_TENSOR.ATTN_OUT,
+ MODEL_TENSOR.FFN_GATE,
+ MODEL_TENSOR.FFN_DOWN,
+ MODEL_TENSOR.FFN_UP,
+ MODEL_TENSOR.ATTN_NORM,
+ MODEL_TENSOR.ATTN_POST_NORM,
+ MODEL_TENSOR.FFN_PRE_NORM,
+ MODEL_TENSOR.FFN_POST_NORM,
+ ],
MODEL_ARCH.STARCODER2: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
@@ -766,6 +915,19 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
+ MODEL_ARCH.OPENELM: [
+ MODEL_TENSOR.TOKEN_EMBD,
+ MODEL_TENSOR.OUTPUT_NORM,
+ MODEL_TENSOR.ATTN_NORM,
+ MODEL_TENSOR.ATTN_QKV,
+ MODEL_TENSOR.ATTN_Q_NORM,
+ MODEL_TENSOR.ATTN_K_NORM,
+ MODEL_TENSOR.ATTN_OUT,
+ MODEL_TENSOR.FFN_NORM,
+ MODEL_TENSOR.FFN_GATE,
+ MODEL_TENSOR.FFN_DOWN,
+ MODEL_TENSOR.FFN_UP,
+ ],
MODEL_ARCH.ARCTIC: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
@@ -814,17 +976,26 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN_SHEXP,
MODEL_TENSOR.FFN_UP_SHEXP,
],
+ MODEL_ARCH.CHATGLM : [
+ MODEL_TENSOR.TOKEN_EMBD,
+ MODEL_TENSOR.ROPE_FREQS,
+ MODEL_TENSOR.OUTPUT_NORM,
+ MODEL_TENSOR.OUTPUT,
+ MODEL_TENSOR.ATTN_NORM,
+ MODEL_TENSOR.ATTN_QKV,
+ MODEL_TENSOR.ATTN_OUT,
+ MODEL_TENSOR.FFN_NORM,
+ MODEL_TENSOR.FFN_DOWN,
+ MODEL_TENSOR.FFN_UP,
+ ],
MODEL_ARCH.BITNET: [
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
- MODEL_TENSOR.ROPE_FREQS,
MODEL_TENSOR.ATTN_NORM,
- MODEL_TENSOR.ATTN_QKV,
MODEL_TENSOR.ATTN_OUT,
- MODEL_TENSOR.ATTN_ROT_EMBD,
MODEL_TENSOR.FFN_NORM,
MODEL_TENSOR.FFN_GATE,
MODEL_TENSOR.FFN_DOWN,
@@ -832,6 +1003,50 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.ATTN_SUB_NORM,
MODEL_TENSOR.FFN_SUB_NORM,
],
+ MODEL_ARCH.T5: [
+ MODEL_TENSOR.TOKEN_EMBD,
+ MODEL_TENSOR.OUTPUT,
+ MODEL_TENSOR.DEC_ATTN_NORM,
+ MODEL_TENSOR.DEC_ATTN_Q,
+ MODEL_TENSOR.DEC_ATTN_K,
+ MODEL_TENSOR.DEC_ATTN_V,
+ MODEL_TENSOR.DEC_ATTN_OUT,
+ MODEL_TENSOR.DEC_ATTN_REL_B,
+ MODEL_TENSOR.DEC_CROSS_ATTN_NORM,
+ MODEL_TENSOR.DEC_CROSS_ATTN_Q,
+ MODEL_TENSOR.DEC_CROSS_ATTN_K,
+ MODEL_TENSOR.DEC_CROSS_ATTN_V,
+ MODEL_TENSOR.DEC_CROSS_ATTN_OUT,
+ MODEL_TENSOR.DEC_CROSS_ATTN_REL_B,
+ MODEL_TENSOR.DEC_FFN_NORM,
+ MODEL_TENSOR.DEC_FFN_GATE,
+ MODEL_TENSOR.DEC_FFN_DOWN,
+ MODEL_TENSOR.DEC_FFN_UP,
+ MODEL_TENSOR.DEC_OUTPUT_NORM,
+ MODEL_TENSOR.ENC_ATTN_NORM,
+ MODEL_TENSOR.ENC_ATTN_Q,
+ MODEL_TENSOR.ENC_ATTN_K,
+ MODEL_TENSOR.ENC_ATTN_V,
+ MODEL_TENSOR.ENC_ATTN_OUT,
+ MODEL_TENSOR.ENC_ATTN_REL_B,
+ MODEL_TENSOR.ENC_FFN_NORM,
+ MODEL_TENSOR.ENC_FFN_GATE,
+ MODEL_TENSOR.ENC_FFN_DOWN,
+ MODEL_TENSOR.ENC_FFN_UP,
+ MODEL_TENSOR.ENC_OUTPUT_NORM,
+ ],
+ MODEL_ARCH.JAIS: [
+ MODEL_TENSOR.TOKEN_EMBD,
+ MODEL_TENSOR.OUTPUT_NORM,
+ MODEL_TENSOR.OUTPUT,
+ MODEL_TENSOR.ATTN_NORM,
+ MODEL_TENSOR.ATTN_QKV,
+ MODEL_TENSOR.ATTN_OUT,
+ MODEL_TENSOR.FFN_NORM,
+ MODEL_TENSOR.FFN_DOWN,
+ MODEL_TENSOR.FFN_GATE,
+ MODEL_TENSOR.FFN_UP,
+ ],
# TODO
}
@@ -869,6 +1084,9 @@ MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.ROPE_FREQS,
MODEL_TENSOR.ATTN_ROT_EMBD,
],
+ MODEL_ARCH.CHATGLM: [
+ MODEL_TENSOR.ROPE_FREQS,
+ ],
}
#
@@ -1056,7 +1274,6 @@ KEY_GENERAL_URL = Keys.General.URL
KEY_GENERAL_DESCRIPTION = Keys.General.DESCRIPTION
KEY_GENERAL_LICENSE = Keys.General.LICENSE
KEY_GENERAL_SOURCE_URL = Keys.General.SOURCE_URL
-KEY_GENERAL_SOURCE_HF_REPO = Keys.General.SOURCE_HF_REPO
KEY_GENERAL_FILE_TYPE = Keys.General.FILE_TYPE
# LLM
diff --git a/gguf-py/gguf/gguf_reader.py b/gguf-py/gguf/gguf_reader.py
index e48bc00c..e8e61abf 100644
--- a/gguf-py/gguf/gguf_reader.py
+++ b/gguf-py/gguf/gguf_reader.py
@@ -67,8 +67,9 @@ class ReaderTensor(NamedTuple):
class GGUFReader:
# I - same as host, S - swapped
- byte_order: Literal['I'] | Literal['S'] = 'I'
+ byte_order: Literal['I', 'S'] = 'I'
alignment: int = GGUF_DEFAULT_ALIGNMENT
+ data_offset: int
# Note: Internal helper, API may change.
gguf_scalar_to_np: dict[GGUFValueType, type[np.generic]] = {
@@ -85,12 +86,16 @@ class GGUFReader:
GGUFValueType.BOOL: np.bool_,
}
- def __init__(self, path: os.PathLike[str] | str, mode: Literal['r'] | Literal['r+'] | Literal['c'] = 'r'):
+ def __init__(self, path: os.PathLike[str] | str, mode: Literal['r', 'r+', 'c'] = 'r'):
self.data = np.memmap(path, mode = mode)
offs = 0
+
+ # Check for GGUF magic
if self._get(offs, np.uint32, override_order = '<')[0] != GGUF_MAGIC:
raise ValueError('GGUF magic invalid')
offs += 4
+
+ # Check GGUF version
temp_version = self._get(offs, np.uint32)
if temp_version[0] & 65535 == 0:
# If we get 0 here that means it's (probably) a GGUF file created for
@@ -103,12 +108,16 @@ class GGUFReader:
self.fields: OrderedDict[str, ReaderField] = OrderedDict()
self.tensors: list[ReaderTensor] = []
offs += self._push_field(ReaderField(offs, 'GGUF.version', [temp_version], [0], [GGUFValueType.UINT32]))
+
+ # Check tensor count and kv count
temp_counts = self._get(offs, np.uint64, 2)
offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64]))
offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64]))
tensor_count, kv_count = temp_counts
offs = self._build_fields(offs, kv_count)
- offs, tensors_fields = self._build_tensors_fields(offs, tensor_count)
+
+ # Build Tensor Info Fields
+ offs, tensors_fields = self._build_tensor_info(offs, tensor_count)
new_align = self.fields.get('general.alignment')
if new_align is not None:
if new_align.types != [GGUFValueType.UINT32]:
@@ -117,6 +126,7 @@ class GGUFReader:
padding = offs % self.alignment
if padding != 0:
offs += self.alignment - padding
+ self.data_offset = offs
self._build_tensors(offs, tensors_fields)
_DT = TypeVar('_DT', bound = npt.DTypeLike)
@@ -130,7 +140,7 @@ class GGUFReader:
return self.tensors[idx]
def _get(
- self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I'] | Literal['S'] | Literal['<'] = None,
+ self, offset: int, dtype: npt.DTypeLike, count: int = 1, override_order: None | Literal['I', 'S', '<'] = None,
) -> npt.NDArray[Any]:
count = int(count)
itemsize = int(np.empty([], dtype = dtype).itemsize)
@@ -193,18 +203,29 @@ class GGUFReader:
# We can't deal with this one.
raise ValueError('Unknown/unhandled field type {gtype}')
- def _get_tensor(self, orig_offs: int) -> ReaderField:
+ def _get_tensor_info_field(self, orig_offs: int) -> ReaderField:
offs = orig_offs
+
+ # Get Tensor Name
name_len, name_data = self._get_str(offs)
offs += int(name_len.nbytes + name_data.nbytes)
+
+ # Get Tensor Dimensions Count
n_dims = self._get(offs, np.uint32)
offs += int(n_dims.nbytes)
+
+ # Get Tensor Dimension Array
dims = self._get(offs, np.uint64, n_dims[0])
offs += int(dims.nbytes)
+
+ # Get Tensor Encoding Scheme Type
raw_dtype = self._get(offs, np.uint32)
offs += int(raw_dtype.nbytes)
+
+ # Get Tensor Offset
offset_tensor = self._get(offs, np.uint64)
offs += int(offset_tensor.nbytes)
+
return ReaderField(
orig_offs,
str(bytes(name_data), encoding = 'utf-8'),
@@ -233,10 +254,10 @@ class GGUFReader:
offs += field_size
return offs
- def _build_tensors_fields(self, offs: int, count: int) -> tuple[int, list[ReaderField]]:
+ def _build_tensor_info(self, offs: int, count: int) -> tuple[int, list[ReaderField]]:
tensor_fields = []
for _ in range(count):
- field = self._get_tensor(offs)
+ field = self._get_tensor_info_field(offs)
offs += sum(int(part.nbytes) for part in field.parts)
tensor_fields.append(field)
return offs, tensor_fields
diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
index a697f657..ba6f53cd 100644
--- a/gguf-py/gguf/gguf_writer.py
+++ b/gguf-py/gguf/gguf_writer.py
@@ -7,6 +7,8 @@ import struct
import tempfile
from dataclasses import dataclass
from enum import Enum, auto
+from math import prod
+from pathlib import Path
from io import BufferedWriter
from typing import IO, Any, Sequence, Mapping
from string import ascii_letters, digits
@@ -31,6 +33,9 @@ from .quants import quant_shape_from_byte_shape
logger = logging.getLogger(__name__)
+SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf"
+
+
@dataclass
class TensorInfo:
shape: Sequence[int]
@@ -55,11 +60,11 @@ class WriterState(Enum):
class GGUFWriter:
- fout: BufferedWriter | None
- path: os.PathLike[str] | str | None
+ fout: list[BufferedWriter] | None
+ path: Path | None
temp_file: tempfile.SpooledTemporaryFile[bytes] | None
- tensors: dict[str, TensorInfo]
- kv_data: dict[str, GGUFValue]
+ tensors: list[dict[str, TensorInfo]]
+ kv_data: list[dict[str, GGUFValue]]
state: WriterState
_simple_value_packing = {
GGUFValueType.UINT8: "B",
@@ -76,29 +81,89 @@ class GGUFWriter:
}
def __init__(
- self, path: os.PathLike[str] | str | None, arch: str, use_temp_file: bool = False,
- endianess: GGUFEndian = GGUFEndian.LITTLE,
+ self, path: os.PathLike[str] | str | None, arch: str, use_temp_file: bool = False, endianess: GGUFEndian = GGUFEndian.LITTLE,
+ split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False, small_first_shard: bool = False
):
self.fout = None
- self.path = path
+ self.path = Path(path) if path else None
self.arch = arch
self.endianess = endianess
self.data_alignment = GGUF_DEFAULT_ALIGNMENT
self.use_temp_file = use_temp_file
self.temp_file = None
- self.tensors = dict()
- self.kv_data = dict()
+ self.tensors = [{}]
+ self.kv_data = [{}]
+ self.split_max_tensors = split_max_tensors
+ self.split_max_size = split_max_size
+ self.dry_run = dry_run
+ self.small_first_shard = small_first_shard
logger.info("gguf: This GGUF file is for {0} Endian only".format(
"Big" if self.endianess == GGUFEndian.BIG else "Little",
))
self.state = WriterState.NO_FILE
+ if self.small_first_shard:
+ self.tensors.append({})
+
self.add_architecture()
- def open_output_file(self, path: os.PathLike[str] | str | None = None) -> None:
+ def get_total_parameter_count(self) -> tuple[int, int, int, int]:
+ total_params = 0
+ shared_params = 0
+ expert_params = 0
+
+ expert_sum = 0
+ n_expert_tensors = 0
+
+ last_lora_a: tuple[str, TensorInfo] | None = None
+
+ for tensors in self.tensors:
+ for name, info in tensors.items():
+
+ shape = info.shape
+
+ if name.endswith(".lora_a"):
+ last_lora_a = (name, info)
+ continue
+ elif name.endswith(".lora_b"):
+ if last_lora_a is None or last_lora_a[0] != name[:-1] + "a":
+ # Bail when the LoRA pair can't be found trivially
+ logger.warning("can't measure LoRA size correctly, tensor order is unusual")
+ return 0, 0, 0, 0
+ else:
+ shape = (*shape[:-1], last_lora_a[1].shape[-1])
+
+ size = prod(shape)
+
+ if "_exps." in name:
+ expert_params += (size // shape[-3])
+ expert_sum += shape[-3]
+ n_expert_tensors += 1
+ else:
+ shared_params += size
+
+ total_params += size
+
+ # Hopefully this should work even for variable-expert-count models
+ expert_count = (expert_sum // n_expert_tensors) if n_expert_tensors > 0 else 0
+
+ # Negate the total to signal it's likely not exact
+ if last_lora_a is not None:
+ total_params = -total_params
+
+ # NOTE: keep the output in the same order as accepted by 'size_label' in gguf-py/gguf/utility.py
+ return total_params, shared_params, expert_params, expert_count
+
+ def format_shard_names(self, path: Path) -> list[Path]:
+ if len(self.tensors) == 1:
+ return [path]
+ return [path.with_name(SHARD_NAME_FORMAT.format(path.stem, i + 1, len(self.tensors))) for i in range(len(self.tensors))]
+
+ def open_output_file(self, path: Path | None = None) -> None:
if self.state is WriterState.EMPTY and self.fout is not None and (path is None or path == self.path):
# allow calling this multiple times as long as the path is the same
return
+
if self.state is not WriterState.NO_FILE:
raise ValueError(f'Expected output file to be not yet opened, got {self.state}')
@@ -106,22 +171,60 @@ class GGUFWriter:
self.path = path
if self.path is not None:
- if self.fout is not None:
- self.fout.close()
- self.fout = open(self.path, "wb")
+ filenames = self.print_plan()
+ self.fout = [open(filename, "wb") for filename in filenames]
self.state = WriterState.EMPTY
- def write_header_to_file(self, path: os.PathLike[str] | str | None = None) -> None:
+ def print_plan(self) -> list[Path]:
+ logger.info("Writing the following files:")
+ assert self.path is not None
+ filenames = self.format_shard_names(self.path)
+ assert len(filenames) == len(self.tensors)
+ for name, tensors in zip(filenames, self.tensors):
+ logger.info(f"{name}: n_tensors = {len(tensors)}, total_size = {GGUFWriter.format_n_bytes_to_str(sum(ti.nbytes for ti in tensors.values()))}")
+
+ if self.dry_run:
+ logger.info("Dry run, not writing files")
+ for name in filenames:
+ print(name) # noqa: NP100
+ exit()
+
+ return filenames
+
+ def add_shard_kv_data(self) -> None:
+ if len(self.tensors) == 1:
+ return
+
+ total_tensors = sum(len(t) for t in self.tensors)
+ assert self.fout is not None
+ total_splits = len(self.fout)
+ self.kv_data.extend({} for _ in range(len(self.kv_data), total_splits))
+ for i, kv_data in enumerate(self.kv_data):
+ kv_data[Keys.Split.LLM_KV_SPLIT_NO] = GGUFValue(i, GGUFValueType.UINT16)
+ kv_data[Keys.Split.LLM_KV_SPLIT_COUNT] = GGUFValue(total_splits, GGUFValueType.UINT16)
+ kv_data[Keys.Split.LLM_KV_SPLIT_TENSORS_COUNT] = GGUFValue(total_tensors, GGUFValueType.INT32)
+
+ def write_header_to_file(self, path: Path | None = None) -> None:
+ if len(self.tensors) == 1 and (self.split_max_tensors != 0 or self.split_max_size != 0):
+ logger.warning("Model fails split requirements, not splitting")
+
self.open_output_file(path)
if self.state is not WriterState.EMPTY:
raise ValueError(f'Expected output file to be empty, got {self.state}')
- self._write_packed("<I", GGUF_MAGIC, skip_pack_prefix = True)
- self._write_packed("I", GGUF_VERSION)
- self._write_packed("Q", len(self.tensors))
- self._write_packed("Q", len(self.kv_data))
- self.flush()
+ assert self.fout is not None
+ assert len(self.fout) == len(self.tensors)
+ assert len(self.kv_data) == 1
+
+ self.add_shard_kv_data()
+
+ for fout, tensors, kv_data in zip(self.fout, self.tensors, self.kv_data):
+ fout.write(self._pack("<I", GGUF_MAGIC, skip_pack_prefix = True))
+ fout.write(self._pack("I", GGUF_VERSION))
+ fout.write(self._pack("Q", len(tensors)))
+ fout.write(self._pack("Q", len(kv_data)))
+ fout.flush()
self.state = WriterState.HEADER
def write_kv_data_to_file(self) -> None:
@@ -129,13 +232,15 @@ class GGUFWriter:
raise ValueError(f'Expected output file to contain the header, got {self.state}')
assert self.fout is not None
- kv_data = bytearray()
+ for fout, kv_data in zip(self.fout, self.kv_data):
+ kv_bytes = bytearray()
+
+ for key, val in kv_data.items():
+ kv_bytes += self._pack_val(key, GGUFValueType.STRING, add_vtype=False)
+ kv_bytes += self._pack_val(val.value, val.type, add_vtype=True)
- for key, val in self.kv_data.items():
- kv_data += self._pack_val(key, GGUFValueType.STRING, add_vtype=False)
- kv_data += self._pack_val(val.value, val.type, add_vtype=True)
+ fout.write(kv_bytes)
- self.fout.write(kv_data)
self.flush()
self.state = WriterState.KV_DATA
@@ -144,28 +249,29 @@ class GGUFWriter:
raise ValueError(f'Expected output file to contain KV data, got {self.state}')
assert self.fout is not None
- ti_data = bytearray()
- offset_tensor = 0
-
- for name, ti in self.tensors.items():
- ti_data += self._pack_val(name, GGUFValueType.STRING, add_vtype=False)
- n_dims = len(ti.shape)
- ti_data += self._pack("I", n_dims)
- for i in range(n_dims):
- ti_data += self._pack("Q", ti.shape[n_dims - 1 - i])
- ti_data += self._pack("I", ti.dtype)
- ti_data += self._pack("Q", offset_tensor)
- offset_tensor += GGUFWriter.ggml_pad(ti.nbytes, self.data_alignment)
-
- self.fout.write(ti_data)
- self.flush()
+ for fout, tensors in zip(self.fout, self.tensors):
+ ti_data = bytearray()
+ offset_tensor = 0
+
+ for name, ti in tensors.items():
+ ti_data += self._pack_val(name, GGUFValueType.STRING, add_vtype=False)
+ n_dims = len(ti.shape)
+ ti_data += self._pack("I", n_dims)
+ for j in range(n_dims):
+ ti_data += self._pack("Q", ti.shape[n_dims - 1 - j])
+ ti_data += self._pack("I", ti.dtype)
+ ti_data += self._pack("Q", offset_tensor)
+ offset_tensor += GGUFWriter.ggml_pad(ti.nbytes, self.data_alignment)
+
+ fout.write(ti_data)
+ fout.flush()
self.state = WriterState.TI_DATA
def add_key_value(self, key: str, val: Any, vtype: GGUFValueType) -> None:
- if key in self.kv_data:
+ if any(key in kv_data for kv_data in self.kv_data):
raise ValueError(f'Duplicated key name {key!r}')
- self.kv_data[key] = GGUFValue(value=val, type=vtype)
+ self.kv_data[0][key] = GGUFValue(value=val, type=vtype)
def add_uint8(self, key: str, val: int) -> None:
self.add_key_value(key,val, GGUFValueType.UINT8)
@@ -206,9 +312,6 @@ class GGUFWriter:
self.add_key_value(key, val, GGUFValueType.STRING)
def add_array(self, key: str, val: Sequence[Any]) -> None:
- if not isinstance(val, Sequence):
- raise ValueError("Value must be a sequence for array type")
-
self.add_key_value(key, val, GGUFValueType.ARRAY)
@staticmethod
@@ -222,7 +325,7 @@ class GGUFWriter:
if self.state is not WriterState.NO_FILE:
raise ValueError(f'Expected output file to be not yet opened, got {self.state}')
- if name in self.tensors:
+ if any(name in tensors for tensors in self.tensors):
raise ValueError(f'Duplicated tensor name {name!r}')
if raw_dtype is None:
@@ -247,7 +350,18 @@ class GGUFWriter:
if tensor_dtype == np.uint8:
tensor_shape = quant_shape_from_byte_shape(tensor_shape, raw_dtype)
- self.tensors[name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes)
+ # make sure there is at least one tensor before splitting
+ if len(self.tensors[-1]) > 0:
+ if ( # split when over tensor limit
+ self.split_max_tensors != 0
+ and len(self.tensors[-1]) >= self.split_max_tensors
+ ) or ( # split when over size limit
+ self.split_max_size != 0
+ and sum(ti.nbytes for ti in self.tensors[-1].values()) + tensor_nbytes > self.split_max_size
+ ):
+ self.tensors.append({})
+
+ self.tensors[-1][name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes)
def add_tensor(
self, name: str, tensor: np.ndarray[Any, Any], raw_shape: Sequence[int] | None = None,
@@ -264,7 +378,7 @@ class GGUFWriter:
self.add_tensor_info(name, shape, tensor.dtype, tensor.nbytes, raw_dtype=raw_dtype)
if self.temp_file is None:
- self.tensors[name].tensor = tensor
+ self.tensors[-1][name].tensor = tensor
return
tensor.tofile(self.temp_file)
@@ -282,9 +396,24 @@ class GGUFWriter:
if self.endianess == GGUFEndian.BIG:
tensor.byteswap(inplace=True)
- self.write_padding(self.fout, self.fout.tell())
- tensor.tofile(self.fout)
- self.write_padding(self.fout, tensor.nbytes)
+
+ file_id = -1
+ for i, tensors in enumerate(self.tensors):
+ if len(tensors) > 0:
+ file_id = i
+ break
+
+ fout = self.fout[file_id]
+
+ # pop the first tensor info
+ # TODO: cleaner way to get the first key
+ first_tensor_name = [name for name, _ in zip(self.tensors[file_id].keys(), range(1))][0]
+ ti = self.tensors[file_id].pop(first_tensor_name)
+ assert ti.nbytes == tensor.nbytes
+
+ self.write_padding(fout, fout.tell())
+ tensor.tofile(fout)
+ self.write_padding(fout, tensor.nbytes)
self.state = WriterState.WEIGHTS
@@ -293,31 +422,43 @@ class GGUFWriter:
assert self.fout is not None
- self.write_padding(self.fout, self.fout.tell())
+ for fout in self.fout:
+ self.write_padding(fout, fout.tell())
if self.temp_file is None:
+ shard_bar = None
bar = None
if progress:
from tqdm import tqdm
- total_bytes = sum(t.nbytes for t in self.tensors.values())
+ total_bytes = sum(ti.nbytes for t in self.tensors for ti in t.values())
+ if len(self.fout) > 1:
+ shard_bar = tqdm(desc=f"Shard (0/{len(self.fout)})", total=None, unit="byte", unit_scale=True)
bar = tqdm(desc="Writing", total=total_bytes, unit="byte", unit_scale=True)
- # relying on the fact that Python dicts preserve insertion order (since 3.7)
- for ti in self.tensors.values():
- assert ti.tensor is not None # can only iterate once over the tensors
- assert ti.tensor.nbytes == ti.nbytes
- ti.tensor.tofile(self.fout)
- if bar is not None:
- bar.update(ti.nbytes)
- self.write_padding(self.fout, ti.nbytes)
- ti.tensor = None
+ for i, (fout, tensors) in enumerate(zip(self.fout, self.tensors)):
+ if shard_bar is not None:
+ shard_bar.set_description(f"Shard ({i + 1}/{len(self.fout)})")
+ total = sum(ti.nbytes for ti in tensors.values())
+ shard_bar.reset(total=(total if total > 0 else None))
+
+ # relying on the fact that Python dicts preserve insertion order (since 3.7)
+ for ti in tensors.values():
+ assert ti.tensor is not None # can only iterate once over the tensors
+ assert ti.tensor.nbytes == ti.nbytes
+ ti.tensor.tofile(fout)
+ if shard_bar is not None:
+ shard_bar.update(ti.nbytes)
+ if bar is not None:
+ bar.update(ti.nbytes)
+ self.write_padding(fout, ti.nbytes)
+ ti.tensor = None
else:
self.temp_file.seek(0)
- shutil.copyfileobj(self.temp_file, self.fout)
+ shutil.copyfileobj(self.temp_file, self.fout[0 if not self.small_first_shard else 1])
self.flush()
self.temp_file.close()
@@ -325,53 +466,129 @@ class GGUFWriter:
def flush(self) -> None:
assert self.fout is not None
- self.fout.flush()
+ for fout in self.fout:
+ fout.flush()
def close(self) -> None:
if self.fout is not None:
- self.fout.close()
+ for fout in self.fout:
+ fout.close()
self.fout = None
+ def add_type(self, type_name: str) -> None:
+ self.add_string(Keys.General.TYPE, type_name)
+
def add_architecture(self) -> None:
self.add_string(Keys.General.ARCHITECTURE, self.arch)
+ def add_quantization_version(self, quantization_version: int) -> None:
+ self.add_uint32(Keys.General.QUANTIZATION_VERSION, quantization_version)
+
+ def add_custom_alignment(self, alignment: int) -> None:
+ self.data_alignment = alignment
+ self.add_uint32(Keys.General.ALIGNMENT, alignment)
+
+ def add_file_type(self, ftype: int) -> None:
+ self.add_uint32(Keys.General.FILE_TYPE, ftype)
+
+ def add_name(self, name: str) -> None:
+ self.add_string(Keys.General.NAME, name)
+
def add_author(self, author: str) -> None:
self.add_string(Keys.General.AUTHOR, author)
def add_version(self, version: str) -> None:
self.add_string(Keys.General.VERSION, version)
- def add_tensor_data_layout(self, layout: str) -> None:
- self.add_string(Keys.LLM.TENSOR_DATA_LAYOUT.format(arch=self.arch), layout)
+ def add_organization(self, organization: str) -> None:
+ self.add_string(Keys.General.ORGANIZATION, organization)
- def add_url(self, url: str) -> None:
- self.add_string(Keys.General.URL, url)
+ def add_finetune(self, finetune: str) -> None:
+ self.add_string(Keys.General.FINETUNE, finetune)
+
+ def add_basename(self, basename: str) -> None:
+ self.add_string(Keys.General.BASENAME, basename)
def add_description(self, description: str) -> None:
self.add_string(Keys.General.DESCRIPTION, description)
- def add_licence(self, licence: str) -> None:
- self.add_string(Keys.General.LICENSE, licence)
+ def add_quantized_by(self, quantized: str) -> None:
+ self.add_string(Keys.General.QUANTIZED_BY, quantized)
+
+ def add_size_label(self, size_label: str) -> None:
+ self.add_string(Keys.General.SIZE_LABEL, size_label)
+
+ def add_license(self, license: str) -> None:
+ self.add_string(Keys.General.LICENSE, license)
+
+ def add_license_name(self, license: str) -> None:
+ self.add_string(Keys.General.LICENSE_NAME, license)
+
+ def add_license_link(self, license: str) -> None:
+ self.add_string(Keys.General.LICENSE_LINK, license)
+
+ def add_url(self, url: str) -> None:
+ self.add_string(Keys.General.URL, url)
+
+ def add_doi(self, doi: str) -> None:
+ self.add_string(Keys.General.DOI, doi)
+
+ def add_uuid(self, uuid: str) -> None:
+ self.add_string(Keys.General.UUID, uuid)
+
+ def add_repo_url(self, repo_url: str) -> None:
+ self.add_string(Keys.General.REPO_URL, repo_url)
def add_source_url(self, url: str) -> None:
self.add_string(Keys.General.SOURCE_URL, url)
- def add_source_hf_repo(self, repo: str) -> None:
- self.add_string(Keys.General.SOURCE_HF_REPO, repo)
+ def add_source_doi(self, doi: str) -> None:
+ self.add_string(Keys.General.SOURCE_DOI, doi)
- def add_file_type(self, ftype: int) -> None:
- self.add_uint32(Keys.General.FILE_TYPE, ftype)
+ def add_source_uuid(self, uuid: str) -> None:
+ self.add_string(Keys.General.SOURCE_UUID, uuid)
- def add_name(self, name: str) -> None:
- self.add_string(Keys.General.NAME, name)
+ def add_source_repo_url(self, repo_url: str) -> None:
+ self.add_string(Keys.General.SOURCE_REPO_URL, repo_url)
- def add_quantization_version(self, quantization_version: int) -> None:
- self.add_uint32(
- Keys.General.QUANTIZATION_VERSION, quantization_version)
+ def add_base_model_count(self, source_count: int) -> None:
+ self.add_uint32(Keys.General.BASE_MODEL_COUNT, source_count)
- def add_custom_alignment(self, alignment: int) -> None:
- self.data_alignment = alignment
- self.add_uint32(Keys.General.ALIGNMENT, alignment)
+ def add_base_model_name(self, source_id: int, name: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_NAME.format(id=source_id), name)
+
+ def add_base_model_author(self, source_id: int, author: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_AUTHOR.format(id=source_id), author)
+
+ def add_base_model_version(self, source_id: int, version: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_VERSION.format(id=source_id), version)
+
+ def add_base_model_organization(self, source_id: int, organization: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_ORGANIZATION.format(id=source_id), organization)
+
+ def add_base_model_url(self, source_id: int, url: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_URL.format(id=source_id), url)
+
+ def add_base_model_doi(self, source_id: int, doi: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_DOI.format(id=source_id), doi)
+
+ def add_base_model_uuid(self, source_id: int, uuid: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_UUID.format(id=source_id), uuid)
+
+ def add_base_model_repo_url(self, source_id: int, repo_url: str) -> None:
+ self.add_string(Keys.General.BASE_MODEL_REPO_URL.format(id=source_id), repo_url)
+
+ def add_tags(self, tags: Sequence[str]) -> None:
+ self.add_array(Keys.General.TAGS, tags)
+
+ def add_languages(self, languages: Sequence[str]) -> None:
+ self.add_array(Keys.General.LANGUAGES, languages)
+
+ def add_datasets(self, datasets: Sequence[str]) -> None:
+ self.add_array(Keys.General.DATASETS, datasets)
+
+ def add_tensor_data_layout(self, layout: str) -> None:
+ self.add_string(Keys.LLM.TENSOR_DATA_LAYOUT.format(arch=self.arch), layout)
def add_vocab_size(self, size: int) -> None:
self.add_uint32(Keys.LLM.VOCAB_SIZE.format(arch=self.arch), size)
@@ -388,8 +605,11 @@ class GGUFWriter:
def add_leading_dense_block_count(self, length: int) -> None:
self.add_uint32(Keys.LLM.LEADING_DENSE_BLOCK_COUNT.format(arch=self.arch), length)
- def add_feed_forward_length(self, length: int) -> None:
- self.add_uint32(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
+ def add_feed_forward_length(self, length: int | Sequence[int]) -> None:
+ if isinstance(length, int):
+ self.add_uint32(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
+ else:
+ self.add_array(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
def add_expert_feed_forward_length(self, length: int) -> None:
self.add_uint32(Keys.LLM.EXPERT_FEED_FORWARD_LENGTH.format(arch=self.arch), length)
@@ -400,11 +620,20 @@ class GGUFWriter:
def add_parallel_residual(self, use: bool) -> None:
self.add_bool(Keys.LLM.USE_PARALLEL_RESIDUAL.format(arch=self.arch), use)
- def add_head_count(self, count: int) -> None:
- self.add_uint32(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count)
+ def add_decoder_start_token_id(self, id: int) -> None:
+ self.add_uint32(Keys.LLM.DECODER_START_TOKEN_ID.format(arch=self.arch), id)
- def add_head_count_kv(self, count: int) -> None:
- self.add_uint32(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count)
+ def add_head_count(self, count: int | Sequence[int]) -> None:
+ if isinstance(count, int):
+ self.add_uint32(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count)
+ else:
+ self.add_array(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count)
+
+ def add_head_count_kv(self, count: int | Sequence[int]) -> None:
+ if isinstance(count, int):
+ self.add_uint32(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count)
+ else:
+ self.add_array(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count)
def add_key_length(self, length: int) -> None:
self.add_uint32(Keys.Attention.KEY_LENGTH.format(arch=self.arch), length)
@@ -421,6 +650,12 @@ class GGUFWriter:
def add_logit_scale(self, value: float) -> None:
self.add_float32(Keys.LLM.LOGIT_SCALE.format(arch=self.arch), value)
+ def add_attn_logit_softcapping(self, value: float) -> None:
+ self.add_float32(Keys.LLM.ATTN_LOGIT_SOFTCAPPING.format(arch=self.arch), value)
+
+ def add_final_logit_softcapping(self, value: float) -> None:
+ self.add_float32(Keys.LLM.FINAL_LOGIT_SOFTCAPPING.format(arch=self.arch), value)
+
def add_expert_count(self, count: int) -> None:
self.add_uint32(Keys.LLM.EXPERT_COUNT.format(arch=self.arch), count)
@@ -448,6 +683,12 @@ class GGUFWriter:
def add_kv_lora_rank(self, length: int) -> None:
self.add_uint32(Keys.Attention.KV_LORA_RANK.format(arch=self.arch), length)
+ def add_relative_attn_buckets_count(self, value: int) -> None:
+ self.add_uint32(Keys.Attention.REL_BUCKETS_COUNT.format(arch=self.arch), value)
+
+ def add_sliding_window(self, value: int) -> None:
+ self.add_uint32(Keys.Attention.SLIDING_WINDOW.format(arch=self.arch), value)
+
def add_pooling_type(self, value: PoolingType) -> None:
self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value.value)
@@ -538,6 +779,12 @@ class GGUFWriter:
def add_add_space_prefix(self, value: bool) -> None:
self.add_bool(Keys.Tokenizer.ADD_PREFIX, value)
+ def add_remove_extra_whitespaces(self, value: bool) -> None:
+ self.add_bool(Keys.Tokenizer.REMOVE_EXTRA_WS, value)
+
+ def add_precompiled_charsmap(self, charsmap: Sequence[bytes]) -> None:
+ self.add_array(Keys.Tokenizer.PRECOMPILED_CHARSMAP, charsmap)
+
def add_chat_template(self, value: str | Sequence[Mapping[str, str]]) -> None:
if not isinstance(value, str):
template_default = None
@@ -599,9 +846,12 @@ class GGUFWriter:
kv_data += self._pack("Q", len(encoded_val))
kv_data += encoded_val
elif vtype == GGUFValueType.ARRAY and isinstance(val, Sequence) and val:
- ltype = GGUFValueType.get_type(val[0])
- if not all(GGUFValueType.get_type(i) is ltype for i in val[1:]):
- raise ValueError("All items in a GGUF array should be of the same type")
+ if isinstance(val, bytes):
+ ltype = GGUFValueType.UINT8
+ else:
+ ltype = GGUFValueType.get_type(val[0])
+ if not all(GGUFValueType.get_type(i) is ltype for i in val[1:]):
+ raise ValueError("All items in a GGUF array should be of the same type")
kv_data += self._pack("I", ltype)
kv_data += self._pack("Q", len(val))
for item in val:
@@ -611,6 +861,13 @@ class GGUFWriter:
return kv_data
- def _write_packed(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> None:
- assert self.fout is not None
- self.fout.write(self._pack(fmt, value, skip_pack_prefix))
+ @staticmethod
+ def format_n_bytes_to_str(num: int) -> str:
+ if num == 0:
+ return "negligible - metadata only"
+ fnum = float(num)
+ for unit in ("", "K", "M", "G"):
+ if abs(fnum) < 1000.0:
+ return f"{fnum:3.1f}{unit}"
+ fnum /= 1000.0
+ return f"{fnum:.1f}T - over 1TB, split recommended"
diff --git a/gguf-py/gguf/lazy.py b/gguf-py/gguf/lazy.py
index 1167335b..ac98d9a9 100644
--- a/gguf-py/gguf/lazy.py
+++ b/gguf-py/gguf/lazy.py
@@ -3,10 +3,8 @@ from abc import ABC, ABCMeta, abstractmethod
import logging
from typing import Any, Callable
-from collections import deque
import numpy as np
-from numpy._typing import _Shape
from numpy.typing import DTypeLike
@@ -16,16 +14,16 @@ logger = logging.getLogger(__name__)
class LazyMeta(ABCMeta):
def __new__(cls, name: str, bases: tuple[type, ...], namespace: dict[str, Any], **kwargs):
- def __getattr__(self, __name: str) -> Any:
- meta_attr = getattr(self._meta, __name)
+ def __getattr__(self, name: str) -> Any:
+ meta_attr = getattr(self._meta, name)
if callable(meta_attr):
return type(self)._wrap_fn(
- (lambda s, *args, **kwargs: getattr(s, __name)(*args, **kwargs)),
+ (lambda s, *args, **kwargs: getattr(s, name)(*args, **kwargs)),
use_self=self,
)
elif isinstance(meta_attr, self._tensor_type):
# e.g. self.T with torch.Tensor should still be wrapped
- return type(self)._wrap_fn(lambda s: getattr(s, __name))(self)
+ return type(self)._wrap_fn(lambda s: getattr(s, name))(self)
else:
# no need to wrap non-tensor properties,
# and they likely don't depend on the actual contents of the tensor
@@ -75,20 +73,18 @@ class LazyBase(ABC, metaclass=LazyMeta):
_tensor_type: type
_meta: Any
_data: Any | None
- _lazy: deque[LazyBase] # shared within a graph, to avoid deep recursion when making eager
_args: tuple
- _func: Callable[[tuple], Any] | None
+ _kwargs: dict[str, Any]
+ _func: Callable[[Any], Any] | None
- def __init__(self, *, meta: Any, data: Any | None = None, lazy: deque[LazyBase] | None = None, args: tuple = (), func: Callable[[tuple], Any] | None = None):
+ def __init__(self, *, meta: Any, data: Any | None = None, args: tuple = (), kwargs: dict[str, Any] | None = None, func: Callable[[Any], Any] | None = None):
super().__init__()
self._meta = meta
self._data = data
- self._lazy = lazy if lazy is not None else deque()
self._args = args
+ self._kwargs = kwargs if kwargs is not None else {}
self._func = func
assert self._func is not None or self._data is not None
- if self._data is None:
- self._lazy.append(self)
def __init_subclass__(cls) -> None:
if "_tensor_type" not in cls.__dict__:
@@ -118,6 +114,7 @@ class LazyBase(ABC, metaclass=LazyMeta):
args = ((use_self,) if use_self is not None else ()) + args
meta_args = LazyBase._recurse_apply(args, lambda t: t._meta)
+ # TODO: maybe handle tensors in kwargs too
if isinstance(meta_noop, bool) and not meta_noop:
try:
@@ -141,21 +138,7 @@ class LazyBase(ABC, metaclass=LazyMeta):
res = cls.meta_with_dtype_and_shape(meta_noop, res.shape)
if isinstance(res, cls._tensor_type):
- def collect_replace(t: LazyBase):
- if collect_replace.shared_lazy is None:
- collect_replace.shared_lazy = t._lazy
- else:
- collect_replace.shared_lazy.extend(t._lazy)
- t._lazy = collect_replace.shared_lazy
-
- # emulating a static variable
- collect_replace.shared_lazy = None
-
- LazyBase._recurse_apply(args, collect_replace)
-
- shared_lazy = collect_replace.shared_lazy
-
- return cls(meta=cls.eager_to_meta(res), lazy=shared_lazy, args=args, func=lambda a: fn(*a, **kwargs))
+ return cls(meta=cls.eager_to_meta(res), args=args, kwargs=kwargs, func=fn)
else:
del res # not needed
# non-tensor return likely relies on the contents of the args
@@ -167,25 +150,18 @@ class LazyBase(ABC, metaclass=LazyMeta):
@classmethod
def to_eager(cls, t: Any) -> Any:
def simple_to_eager(_t: LazyBase) -> Any:
- def already_eager_to_eager(_t: LazyBase) -> Any:
- assert _t._data is not None
+ if _t._data is not None:
return _t._data
- while _t._data is None:
- lt = _t._lazy.popleft()
- if lt._data is not None:
- # Lazy tensor did not belong in the lazy queue.
- # Weirdly only happens with Bloom models...
- # likely because tensors aren't unique in the queue.
- # The final output is still the same as in eager mode,
- # so it's safe to ignore this.
- continue
- assert lt._func is not None
- lt._args = cls._recurse_apply(lt._args, already_eager_to_eager)
- lt._data = lt._func(lt._args)
- # sanity check
- assert lt._data.dtype == lt._meta.dtype
- assert lt._data.shape == lt._meta.shape
+ # NOTE: there's a recursion limit in Python (usually 1000)
+
+ assert _t._func is not None
+ _t._args = cls._recurse_apply(_t._args, simple_to_eager)
+ _t._data = _t._func(*_t._args, **_t._kwargs)
+ # sanity check
+ assert _t._data is not None
+ assert _t._data.dtype == _t._meta.dtype
+ assert _t._data.shape == _t._meta.shape
return _t._data
@@ -204,7 +180,7 @@ class LazyBase(ABC, metaclass=LazyMeta):
@classmethod
def from_eager(cls, t: Any) -> Any:
if type(t) is cls:
- # already eager
+ # already lazy
return t
elif isinstance(t, cls._tensor_type):
return cls(meta=cls.eager_to_meta(t), data=t)
@@ -216,7 +192,7 @@ class LazyNumpyTensor(LazyBase):
_tensor_type = np.ndarray
@classmethod
- def meta_with_dtype_and_shape(cls, dtype: DTypeLike, shape: _Shape) -> np.ndarray[Any, Any]:
+ def meta_with_dtype_and_shape(cls, dtype: DTypeLike, shape: tuple[int, ...]) -> np.ndarray[Any, Any]:
# The initial idea was to use np.nan as the fill value,
# but non-float types like np.int16 can't use that.
# So zero it is.
@@ -226,8 +202,7 @@ class LazyNumpyTensor(LazyBase):
def astype(self, dtype, *args, **kwargs):
meta = type(self).meta_with_dtype_and_shape(dtype, self._meta.shape)
full_args = (self, dtype,) + args
- # very important to pass the shared _lazy deque, or else there's an infinite loop somewhere.
- return type(self)(meta=meta, args=full_args, lazy=self._lazy, func=(lambda a: a[0].astype(*a[1:], **kwargs)))
+ return type(self)(meta=meta, args=full_args, kwargs=kwargs, func=(lambda a, *args, **kwargs: a.astype(*args, **kwargs)))
def tofile(self, *args, **kwargs):
eager = LazyNumpyTensor.to_eager(self)
diff --git a/gguf-py/gguf/metadata.py b/gguf-py/gguf/metadata.py
new file mode 100644
index 00000000..15189f71
--- /dev/null
+++ b/gguf-py/gguf/metadata.py
@@ -0,0 +1,503 @@
+from __future__ import annotations
+
+import re
+import json
+import yaml
+import logging
+from pathlib import Path
+from typing import Any, Literal, Optional
+from dataclasses import dataclass
+
+from .constants import Keys
+
+import gguf
+
+logger = logging.getLogger("metadata")
+
+
+@dataclass
+class Metadata:
+ # Authorship Metadata to be written to GGUF KV Store
+ name: Optional[str] = None
+ author: Optional[str] = None
+ version: Optional[str] = None
+ organization: Optional[str] = None
+ finetune: Optional[str] = None
+ basename: Optional[str] = None
+ description: Optional[str] = None
+ quantized_by: Optional[str] = None
+ size_label: Optional[str] = None
+ url: Optional[str] = None
+ doi: Optional[str] = None
+ uuid: Optional[str] = None
+ repo_url: Optional[str] = None
+ source_url: Optional[str] = None
+ source_doi: Optional[str] = None
+ source_uuid: Optional[str] = None
+ source_repo_url: Optional[str] = None
+ license: Optional[str] = None
+ license_name: Optional[str] = None
+ license_link: Optional[str] = None
+ base_models: Optional[list[dict]] = None
+ tags: Optional[list[str]] = None
+ languages: Optional[list[str]] = None
+ datasets: Optional[list[str]] = None
+
+ @staticmethod
+ def load(metadata_override_path: Optional[Path] = None, model_path: Optional[Path] = None, model_name: Optional[str] = None, total_params: int = 0) -> Metadata:
+ # This grabs as many contextual authorship metadata as possible from the model repository
+ # making any conversion as required to match the gguf kv store metadata format
+ # as well as giving users the ability to override any authorship metadata that may be incorrect
+
+ # Create a new Metadata instance
+ metadata = Metadata()
+
+ model_card = Metadata.load_model_card(model_path)
+ hf_params = Metadata.load_hf_parameters(model_path)
+ # TODO: load adapter_config.json when possible, it usually contains the base model of the LoRA adapter
+
+ # heuristics
+ metadata = Metadata.apply_metadata_heuristic(metadata, model_card, hf_params, model_path, total_params)
+
+ # Metadata Override File Provided
+ # This is based on LLM_KV_NAMES mapping in llama.cpp
+ metadata_override = Metadata.load_metadata_override(metadata_override_path)
+
+ metadata.name = metadata_override.get(Keys.General.NAME, metadata.name)
+ metadata.author = metadata_override.get(Keys.General.AUTHOR, metadata.author)
+ metadata.version = metadata_override.get(Keys.General.VERSION, metadata.version)
+ metadata.organization = metadata_override.get(Keys.General.ORGANIZATION, metadata.organization)
+
+ metadata.finetune = metadata_override.get(Keys.General.FINETUNE, metadata.finetune)
+ metadata.basename = metadata_override.get(Keys.General.BASENAME, metadata.basename)
+
+ metadata.description = metadata_override.get(Keys.General.DESCRIPTION, metadata.description)
+ metadata.quantized_by = metadata_override.get(Keys.General.QUANTIZED_BY, metadata.quantized_by)
+
+ metadata.size_label = metadata_override.get(Keys.General.SIZE_LABEL, metadata.size_label)
+ metadata.license_name = metadata_override.get(Keys.General.LICENSE_NAME, metadata.license_name)
+ metadata.license_link = metadata_override.get(Keys.General.LICENSE_LINK, metadata.license_link)
+
+ metadata.url = metadata_override.get(Keys.General.URL, metadata.url)
+ metadata.doi = metadata_override.get(Keys.General.DOI, metadata.doi)
+ metadata.uuid = metadata_override.get(Keys.General.UUID, metadata.uuid)
+ metadata.repo_url = metadata_override.get(Keys.General.REPO_URL, metadata.repo_url)
+
+ metadata.source_url = metadata_override.get(Keys.General.SOURCE_URL, metadata.source_url)
+ metadata.source_doi = metadata_override.get(Keys.General.SOURCE_DOI, metadata.source_doi)
+ metadata.source_uuid = metadata_override.get(Keys.General.SOURCE_UUID, metadata.source_uuid)
+ metadata.source_repo_url = metadata_override.get(Keys.General.SOURCE_REPO_URL, metadata.source_repo_url)
+
+ # Base Models is received here as an array of models
+ metadata.base_models = metadata_override.get("general.base_models", metadata.base_models)
+
+ metadata.tags = metadata_override.get(Keys.General.TAGS, metadata.tags)
+ metadata.languages = metadata_override.get(Keys.General.LANGUAGES, metadata.languages)
+ metadata.datasets = metadata_override.get(Keys.General.DATASETS, metadata.datasets)
+
+ # Direct Metadata Override (via direct cli argument)
+ if model_name is not None:
+ metadata.name = model_name
+
+ return metadata
+
+ @staticmethod
+ def load_metadata_override(metadata_override_path: Optional[Path] = None) -> dict[str, Any]:
+ if metadata_override_path is None or not metadata_override_path.is_file():
+ return {}
+
+ with open(metadata_override_path, "r", encoding="utf-8") as f:
+ return json.load(f)
+
+ @staticmethod
+ def load_model_card(model_path: Optional[Path] = None) -> dict[str, Any]:
+ if model_path is None or not model_path.is_dir():
+ return {}
+
+ model_card_path = model_path / "README.md"
+
+ if not model_card_path.is_file():
+ return {}
+
+ # The model card metadata is assumed to always be in YAML
+ # ref: https://github.com/huggingface/transformers/blob/a5c642fe7a1f25d3bdcd76991443ba6ff7ee34b2/src/transformers/modelcard.py#L468-L473
+ with open(model_card_path, "r", encoding="utf-8") as f:
+ if f.readline() == "---\n":
+ raw = f.read().partition("---\n")[0]
+ data = yaml.safe_load(raw)
+ if isinstance(data, dict):
+ return data
+ else:
+ logger.error(f"while reading YAML model card frontmatter, data is {type(data)} instead of dict")
+ return {}
+ else:
+ return {}
+
+ @staticmethod
+ def load_hf_parameters(model_path: Optional[Path] = None) -> dict[str, Any]:
+ if model_path is None or not model_path.is_dir():
+ return {}
+
+ config_path = model_path / "config.json"
+
+ if not config_path.is_file():
+ return {}
+
+ with open(config_path, "r", encoding="utf-8") as f:
+ return json.load(f)
+
+ @staticmethod
+ def id_to_title(string):
+ # Convert capitalization into title form unless acronym or version number
+ return ' '.join([w.title() if w.islower() and not re.match(r'^(v\d+(?:\.\d+)*|\d.*)$', w) else w for w in string.strip().replace('-', ' ').split()])
+
+ @staticmethod
+ def get_model_id_components(model_id: Optional[str] = None, total_params: int = 0) -> tuple[str | None, str | None, str | None, str | None, str | None, str | None]:
+ # Huggingface often store model id as '<org>/<model name>'
+ # so let's parse it and apply some heuristics if possible for model name components
+
+ if model_id is None:
+ # model ID missing
+ return None, None, None, None, None, None
+
+ if ' ' in model_id:
+ # model ID is actually a normal human sentence
+ # which means its most likely a normal model name only
+ # not part of the hugging face naming standard, but whatever
+ return model_id, None, None, None, None, None
+
+ if '/' in model_id:
+ # model ID (huggingface style)
+ org_component, model_full_name_component = model_id.split('/', 1)
+ else:
+ # model ID but missing org components
+ org_component, model_full_name_component = None, model_id
+
+ # Check if we erroneously matched against './' or '../' etc...
+ if org_component is not None and org_component[0] == '.':
+ org_component = None
+
+ name_parts: list[str] = model_full_name_component.split('-')
+
+ # Remove empty parts
+ for i in reversed(range(len(name_parts))):
+ if len(name_parts[i]) == 0:
+ del name_parts[i]
+
+ name_types: list[
+ set[Literal["basename", "size_label", "finetune", "version", "type"]]
+ ] = [set() for _ in name_parts]
+
+ # Annotate the name
+ for i, part in enumerate(name_parts):
+ # Version
+ if re.fullmatch(r'(v|iter)?\d+([.]\d+)*', part, re.IGNORECASE):
+ name_types[i].add("version")
+ # Quant type (should not be there for base models, but still annotated)
+ elif re.fullmatch(r'i?q\d(_\w)*|b?fp?(16|32)', part, re.IGNORECASE):
+ name_types[i].add("type")
+ name_parts[i] = part.upper()
+ # Model size
+ elif i > 0 and re.fullmatch(r'(([A]|\d+[x])?\d+([._]\d+)?[KMBT][\d]?|small|mini|medium|large|x?xl)', part, re.IGNORECASE):
+ part = part.replace("_", ".")
+ # Handle weird bloom-7b1 notation
+ if part[-1].isdecimal():
+ part = part[:-2] + "." + part[-1] + part[-2]
+ # Normalize the size suffixes
+ if len(part) > 1 and part[-2].isdecimal():
+ if part[-1] in "kmbt":
+ part = part[:-1] + part[-1].upper()
+ if total_params != 0:
+ try:
+ label_params = float(part[:-1]) * pow(1000, " KMBT".find(part[-1]))
+ # Only use it as a size label if it's close or bigger than the model size
+ # Note that LoRA adapters don't necessarily include all layers,
+ # so this is why bigger label sizes are accepted.
+ # Do not use the size label when it's smaller than 1/8 of the model size
+ if (total_params < 0 and label_params < abs(total_params) // 8) or (
+ # Check both directions when the current model isn't a LoRA adapter
+ total_params > 0 and abs(label_params - total_params) > 7 * total_params // 8
+ ):
+ # Likely a context length
+ name_types[i].add("finetune")
+ # Lowercase the size when it's a context length
+ part = part[:-1] + part[-1].lower()
+ except ValueError:
+ # Failed to convert the size label to float, use it anyway
+ pass
+ if len(name_types[i]) == 0:
+ name_types[i].add("size_label")
+ name_parts[i] = part
+ # Some easy to recognize finetune names
+ elif i > 0 and re.fullmatch(r'chat|instruct|vision|lora', part, re.IGNORECASE):
+ if total_params < 0 and part.lower() == "lora":
+ # ignore redundant "lora" in the finetune part when the output is a lora adapter
+ name_types[i].add("type")
+ else:
+ name_types[i].add("finetune")
+
+ # Ignore word-based size labels when there is at least a number-based one present
+ # TODO: should word-based size labels always be removed instead?
+ if any(c.isdecimal() for n, t in zip(name_parts, name_types) if "size_label" in t for c in n):
+ for n, t in zip(name_parts, name_types):
+ if "size_label" in t:
+ if all(c.isalpha() for c in n):
+ t.remove("size_label")
+
+ at_start = True
+ # Find the basename through the annotated name
+ for part, t in zip(name_parts, name_types):
+ if at_start and ((len(t) == 0 and part[0].isalpha()) or "version" in t):
+ t.add("basename")
+ else:
+ if at_start:
+ at_start = False
+ if len(t) == 0:
+ t.add("finetune")
+
+ # Remove the basename annotation from trailing version
+ for part, t in zip(reversed(name_parts), reversed(name_types)):
+ if "basename" in t and len(t) > 1:
+ t.remove("basename")
+ else:
+ break
+
+ basename = "-".join(n for n, t in zip(name_parts, name_types) if "basename" in t) or None
+ # Deduplicate size labels using order-preserving 'dict' ('set' seems to sort the keys)
+ size_label = "-".join(dict.fromkeys(s for s, t in zip(name_parts, name_types) if "size_label" in t).keys()) or None
+ finetune = "-".join(f for f, t in zip(name_parts, name_types) if "finetune" in t) or None
+ # TODO: should the basename version always be excluded?
+ # NOTE: multiple finetune versions are joined together
+ version = "-".join(v for v, t, in zip(name_parts, name_types) if "version" in t and "basename" not in t) or None
+
+ if size_label is None and finetune is None and version is None:
+ # Too ambiguous, output nothing
+ basename = None
+
+ return model_full_name_component, org_component, basename, finetune, version, size_label
+
+ @staticmethod
+ def apply_metadata_heuristic(metadata: Metadata, model_card: Optional[dict] = None, hf_params: Optional[dict] = None, model_path: Optional[Path] = None, total_params: int = 0) -> Metadata:
+ # Reference Model Card Metadata: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
+
+ # Model Card Heuristics
+ ########################
+ if model_card is not None:
+
+ if "model_name" in model_card and metadata.name is None:
+ # Not part of huggingface model card standard but notice some model creator using it
+ # such as TheBloke in 'TheBloke/Mistral-7B-Instruct-v0.2-GGUF'
+ metadata.name = model_card.get("model_name")
+
+ if "model_creator" in model_card and metadata.author is None:
+ # Not part of huggingface model card standard but notice some model creator using it
+ # such as TheBloke in 'TheBloke/Mistral-7B-Instruct-v0.2-GGUF'
+ metadata.author = model_card.get("model_creator")
+
+ if "model_type" in model_card and metadata.basename is None:
+ # Not part of huggingface model card standard but notice some model creator using it
+ # such as TheBloke in 'TheBloke/Mistral-7B-Instruct-v0.2-GGUF'
+ metadata.basename = model_card.get("model_type")
+
+ if "base_model" in model_card:
+ # This represents the parent models that this is based on
+ # Example: stabilityai/stable-diffusion-xl-base-1.0. Can also be a list (for merges)
+ # Example of merges: https://huggingface.co/EmbeddedLLM/Mistral-7B-Merge-14-v0.1/blob/main/README.md
+ metadata_base_models = []
+ base_model_value = model_card.get("base_model", None)
+
+ if base_model_value is not None:
+ if isinstance(base_model_value, str):
+ metadata_base_models.append(base_model_value)
+ elif isinstance(base_model_value, list):
+ metadata_base_models.extend(base_model_value)
+
+ if metadata.base_models is None:
+ metadata.base_models = []
+
+ for model_id in metadata_base_models:
+ # NOTE: model size of base model is assumed to be similar to the size of the current model
+ model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params)
+ base_model = {}
+ if model_full_name_component is not None:
+ base_model["name"] = Metadata.id_to_title(model_full_name_component)
+ if org_component is not None:
+ base_model["organization"] = Metadata.id_to_title(org_component)
+ if version is not None:
+ base_model["version"] = version
+ if org_component is not None and model_full_name_component is not None:
+ base_model["repo_url"] = f"https://huggingface.co/{org_component}/{model_full_name_component}"
+ metadata.base_models.append(base_model)
+
+ if "license" in model_card and metadata.license is None:
+ metadata.license = model_card.get("license")
+
+ if "license_name" in model_card and metadata.license_name is None:
+ metadata.license_name = model_card.get("license_name")
+
+ if "license_link" in model_card and metadata.license_link is None:
+ metadata.license_link = model_card.get("license_link")
+
+ tags_value = model_card.get("tags", None)
+ if tags_value is not None:
+
+ if metadata.tags is None:
+ metadata.tags = []
+
+ if isinstance(tags_value, str):
+ metadata.tags.append(tags_value)
+ elif isinstance(tags_value, list):
+ metadata.tags.extend(tags_value)
+
+ pipeline_tags_value = model_card.get("pipeline_tag", None)
+ if pipeline_tags_value is not None:
+
+ if metadata.tags is None:
+ metadata.tags = []
+
+ if isinstance(pipeline_tags_value, str):
+ metadata.tags.append(pipeline_tags_value)
+ elif isinstance(pipeline_tags_value, list):
+ metadata.tags.extend(pipeline_tags_value)
+
+ language_value = model_card.get("languages", model_card.get("language", None))
+ if language_value is not None:
+
+ if metadata.languages is None:
+ metadata.languages = []
+
+ if isinstance(language_value, str):
+ metadata.languages.append(language_value)
+ elif isinstance(language_value, list):
+ metadata.languages.extend(language_value)
+
+ dataset_value = model_card.get("datasets", model_card.get("dataset", None))
+ if dataset_value is not None:
+
+ if metadata.datasets is None:
+ metadata.datasets = []
+
+ if isinstance(dataset_value, str):
+ metadata.datasets.append(dataset_value)
+ elif isinstance(dataset_value, list):
+ metadata.datasets.extend(dataset_value)
+
+ # Hugging Face Parameter Heuristics
+ ####################################
+
+ if hf_params is not None:
+
+ hf_name_or_path = hf_params.get("_name_or_path")
+ if hf_name_or_path is not None and hf_name_or_path.count('/') <= 1:
+ # Use _name_or_path only if its actually a model name and not some computer path
+ # e.g. 'meta-llama/Llama-2-7b-hf'
+ model_id = hf_name_or_path
+ model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params)
+ if metadata.name is None and model_full_name_component is not None:
+ metadata.name = Metadata.id_to_title(model_full_name_component)
+ if metadata.organization is None and org_component is not None:
+ metadata.organization = Metadata.id_to_title(org_component)
+ if metadata.basename is None and basename is not None:
+ metadata.basename = basename
+ if metadata.finetune is None and finetune is not None:
+ metadata.finetune = finetune
+ if metadata.version is None and version is not None:
+ metadata.version = version
+ if metadata.size_label is None and size_label is not None:
+ metadata.size_label = size_label
+
+ # Directory Folder Name Fallback Heuristics
+ ############################################
+ if model_path is not None:
+ model_id = model_path.name
+ model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params)
+ if metadata.name is None and model_full_name_component is not None:
+ metadata.name = Metadata.id_to_title(model_full_name_component)
+ if metadata.organization is None and org_component is not None:
+ metadata.organization = Metadata.id_to_title(org_component)
+ if metadata.basename is None and basename is not None:
+ metadata.basename = basename
+ if metadata.finetune is None and finetune is not None:
+ metadata.finetune = finetune
+ if metadata.version is None and version is not None:
+ metadata.version = version
+ if metadata.size_label is None and size_label is not None:
+ metadata.size_label = size_label
+
+ return metadata
+
+ def set_gguf_meta_model(self, gguf_writer: gguf.GGUFWriter):
+ assert self.name is not None
+ gguf_writer.add_name(self.name)
+
+ if self.author is not None:
+ gguf_writer.add_author(self.author)
+ if self.version is not None:
+ gguf_writer.add_version(self.version)
+ if self.organization is not None:
+ gguf_writer.add_organization(self.organization)
+
+ if self.finetune is not None:
+ gguf_writer.add_finetune(self.finetune)
+ if self.basename is not None:
+ gguf_writer.add_basename(self.basename)
+
+ if self.description is not None:
+ gguf_writer.add_description(self.description)
+ if self.quantized_by is not None:
+ gguf_writer.add_quantized_by(self.quantized_by)
+
+ if self.size_label is not None:
+ gguf_writer.add_size_label(self.size_label)
+
+ if self.license is not None:
+ gguf_writer.add_license(self.license)
+ if self.license_name is not None:
+ gguf_writer.add_license_name(self.license_name)
+ if self.license_link is not None:
+ gguf_writer.add_license_link(self.license_link)
+
+ if self.url is not None:
+ gguf_writer.add_url(self.url)
+ if self.doi is not None:
+ gguf_writer.add_doi(self.doi)
+ if self.uuid is not None:
+ gguf_writer.add_uuid(self.uuid)
+ if self.repo_url is not None:
+ gguf_writer.add_repo_url(self.repo_url)
+
+ if self.source_url is not None:
+ gguf_writer.add_source_url(self.source_url)
+ if self.source_doi is not None:
+ gguf_writer.add_source_doi(self.source_doi)
+ if self.source_uuid is not None:
+ gguf_writer.add_source_uuid(self.source_uuid)
+ if self.source_repo_url is not None:
+ gguf_writer.add_source_repo_url(self.source_repo_url)
+
+ if self.base_models is not None:
+ gguf_writer.add_base_model_count(len(self.base_models))
+ for key, base_model_entry in enumerate(self.base_models):
+ if "name" in base_model_entry:
+ gguf_writer.add_base_model_name(key, base_model_entry["name"])
+ if "author" in base_model_entry:
+ gguf_writer.add_base_model_author(key, base_model_entry["author"])
+ if "version" in base_model_entry:
+ gguf_writer.add_base_model_version(key, base_model_entry["version"])
+ if "organization" in base_model_entry:
+ gguf_writer.add_base_model_organization(key, base_model_entry["organization"])
+ if "url" in base_model_entry:
+ gguf_writer.add_base_model_url(key, base_model_entry["url"])
+ if "doi" in base_model_entry:
+ gguf_writer.add_base_model_doi(key, base_model_entry["doi"])
+ if "uuid" in base_model_entry:
+ gguf_writer.add_base_model_uuid(key, base_model_entry["uuid"])
+ if "repo_url" in base_model_entry:
+ gguf_writer.add_base_model_repo_url(key, base_model_entry["repo_url"])
+
+ if self.tags is not None:
+ gguf_writer.add_tags(self.tags)
+ if self.languages is not None:
+ gguf_writer.add_languages(self.languages)
+ if self.datasets is not None:
+ gguf_writer.add_datasets(self.datasets)
diff --git a/gguf-py/gguf/quants.py b/gguf-py/gguf/quants.py
index b22eec16..16e0a9aa 100644
--- a/gguf-py/gguf/quants.py
+++ b/gguf-py/gguf/quants.py
@@ -43,7 +43,7 @@ def __apply_over_grouped_rows(func: Callable[[np.ndarray], np.ndarray], arr: np.
osize *= dim
out = np.empty(shape=osize, dtype=otype)
# compute over groups of 16 rows (arbitrary, but seems good for performance)
- n_groups = rows.shape[0] // 16
+ n_groups = (rows.shape[0] // 16) or 1
np.concatenate([func(group).ravel() for group in np.array_split(rows, n_groups)], axis=0, out=out)
return out.reshape(oshape)
diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py
index 350035bd..9aa2209e 100644
--- a/gguf-py/gguf/tensor_mapping.py
+++ b/gguf-py/gguf/tensor_mapping.py
@@ -10,7 +10,7 @@ class TensorNameMap:
# Token embeddings
MODEL_TENSOR.TOKEN_EMBD: (
"gpt_neox.embed_in", # gptneox
- "transformer.wte", # gpt2 gpt-j mpt refact qwen dbrx
+ "transformer.wte", # gpt2 gpt-j mpt refact qwen dbrx jais
"transformer.word_embeddings", # falcon
"word_embeddings", # bloom
"model.embed_tokens", # llama-hf
@@ -24,6 +24,9 @@ class TensorNameMap:
"backbone.embedding", # mamba
"backbone.embeddings", # mamba-hf
"transformer.in_out_embed", # Grok
+ "embedding.word_embeddings", # chatglm
+ "transformer.token_embeddings", # openelm
+ "shared", # t5
),
# Token type embeddings
@@ -36,6 +39,7 @@ class TensorNameMap:
"word_embeddings_layernorm", # bloom
"embeddings.LayerNorm", # bert
"emb_ln", # nomic-bert
+ "transformer.norm", # openelm
),
# Position embeddings
@@ -48,16 +52,17 @@ class TensorNameMap:
# Output
MODEL_TENSOR.OUTPUT: (
"embed_out", # gptneox
- "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx
+ "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais
"output", # llama-pth bloom internlm2
"word_embeddings_for_head", # persimmon
"lm_head.linear", # phi2
+ "output_layer", # chatglm
),
# Output norm
MODEL_TENSOR.OUTPUT_NORM: (
"gpt_neox.final_layer_norm", # gptneox
- "transformer.ln_f", # gpt2 gpt-j falcon
+ "transformer.ln_f", # gpt2 gpt-j falcon jais
"model.norm", # llama-hf baichuan internlm2
"norm", # llama-pth
"transformer.norm_f", # mpt dbrx
@@ -68,11 +73,14 @@ class TensorNameMap:
"model.norm_f", # mamba-qbert
"backbone.norm_f", # mamba
"transformer.rms_norm", # Grok
+ "encoder.final_layernorm", # chatglm
+ "transformer.norm", # openelm
),
# Rope frequencies
MODEL_TENSOR.ROPE_FREQS: (
"rope.freqs", # llama-pth
+ "rotary_pos_emb.inv_freq", # chatglm
),
}
@@ -80,7 +88,7 @@ class TensorNameMap:
# Attention norm
MODEL_TENSOR.ATTN_NORM: (
"gpt_neox.layers.{bid}.input_layernorm", # gptneox
- "transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen
+ "transformer.h.{bid}.ln_1", # gpt2 gpt-j refact qwen jais
"transformer.blocks.{bid}.norm_1", # mpt
"transformer.h.{bid}.input_layernorm", # falcon7b
"h.{bid}.input_layernorm", # bloom
@@ -97,6 +105,8 @@ class TensorNameMap:
"backbone.layers.{bid}.norm", # mamba
"transformer.decoder_layer.{bid}.rms_norm", # Grok
"transformer.blocks.{bid}.norm_attn_norm.norm_1", # dbrx
+ "encoder.layers.{bid}.input_layernorm", # chatglm
+ "transformer.layers.{bid}.attn_norm", # openelm
),
# Attention norm 2
@@ -108,7 +118,7 @@ class TensorNameMap:
# Attention query-key-value
MODEL_TENSOR.ATTN_QKV: (
"gpt_neox.layers.{bid}.attention.query_key_value", # gptneox
- "transformer.h.{bid}.attn.c_attn", # gpt2 qwen
+ "transformer.h.{bid}.attn.c_attn", # gpt2 qwen jais
"transformer.blocks.{bid}.attn.Wqkv", # mpt
"transformer.blocks.{bid}.norm_attn_norm.attn.Wqkv", # dbrx
"transformer.h.{bid}.self_attention.query_key_value", # falcon
@@ -118,7 +128,9 @@ class TensorNameMap:
"h.{bid}.attn.c_attn", # gpt2
"transformer.h.{bid}.mixer.Wqkv", # phi2
"encoder.layers.{bid}.attn.Wqkv", # nomic-bert
- "model.layers.{bid}.self_attn.qkv_proj" # phi3
+ "model.layers.{bid}.self_attn.qkv_proj", # phi3
+ "encoder.layers.{bid}.self_attention.query_key_value", # chatglm
+ "transformer.layers.{bid}.attn.qkv_proj", # openelm
),
# Attention query
@@ -129,7 +141,7 @@ class TensorNameMap:
"transformer.h.{bid}.attn.q_proj", # gpt-j
"model.layers.layers.{bid}.self_attn.q_proj", # plamo
"model.layers.{bid}.attention.wq", # internlm2
- "transformer.decoder_layer.{bid}.multi_head_attention.query" # Grok
+ "transformer.decoder_layer.{bid}.multi_head_attention.query",# Grok
),
# Attention key
@@ -141,7 +153,7 @@ class TensorNameMap:
"transformer.h.{bid}.attn.k", # refact
"model.layers.layers.{bid}.self_attn.k_proj", # plamo
"model.layers.{bid}.attention.wk", # internlm2
- "transformer.decoder_layer.{bid}.multi_head_attention.key" # Grok
+ "transformer.decoder_layer.{bid}.multi_head_attention.key",# Grok
),
# Attention value
@@ -159,7 +171,7 @@ class TensorNameMap:
# Attention output
MODEL_TENSOR.ATTN_OUT: (
"gpt_neox.layers.{bid}.attention.dense", # gptneox
- "transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen
+ "transformer.h.{bid}.attn.c_proj", # gpt2 refact qwen jais
"transformer.blocks.{bid}.attn.out_proj", # mpt
"transformer.h.{bid}.self_attention.dense", # falcon
"h.{bid}.self_attention.dense", # bloom
@@ -176,6 +188,8 @@ class TensorNameMap:
"encoder.layers.{bid}.attn.out_proj", # nomic-bert
"transformer.decoder_layer.{bid}.multi_head_attention.linear", # Grok
"transformer.blocks.{bid}.norm_attn_norm.attn.out_proj", # dbrx
+ "encoder.layers.{bid}.self_attention.dense", # chatglm
+ "transformer.layers.{bid}.attn.out_proj", # openelm
),
# Attention output norm
@@ -186,6 +200,10 @@ class TensorNameMap:
"transformer.blocks.{bid}.norm_attn_norm.norm_2", # dbrx
),
+ MODEL_TENSOR.ATTN_POST_NORM: (
+ "model.layers.{bid}.post_attention_layernorm", # gemma2
+ ),
+
# Rotary embeddings
MODEL_TENSOR.ATTN_ROT_EMBD: (
"model.layers.{bid}.self_attn.rotary_emb.inv_freq", # llama-hf
@@ -197,7 +215,7 @@ class TensorNameMap:
# Feed-forward norm
MODEL_TENSOR.FFN_NORM: (
"gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox
- "transformer.h.{bid}.ln_2", # gpt2 refact qwen
+ "transformer.h.{bid}.ln_2", # gpt2 refact qwen jais
"h.{bid}.post_attention_layernorm", # bloom
"transformer.blocks.{bid}.norm_2", # mpt
"model.layers.{bid}.post_attention_layernorm", # llama-hf
@@ -207,6 +225,18 @@ class TensorNameMap:
"h.{bid}.ln_2", # gpt2
"model.layers.{bid}.ffn_norm", # internlm2
"transformer.decoder_layer.{bid}.rms_norm_2", # Grok
+ "encoder.layers.{bid}.post_attention_layernorm", # chatglm
+ "transformer.layers.{bid}.ffn_norm", # openelm
+ ),
+
+ # Post feed-forward norm
+ MODEL_TENSOR.FFN_PRE_NORM: (
+ "model.layers.{bid}.pre_feedforward_layernorm", # gemma2
+ ),
+
+ # Post feed-forward norm
+ MODEL_TENSOR.FFN_POST_NORM: (
+ "model.layers.{bid}.post_feedforward_layernorm", # gemma2
),
MODEL_TENSOR.FFN_GATE_INP: (
@@ -224,7 +254,7 @@ class TensorNameMap:
# Feed-forward up
MODEL_TENSOR.FFN_UP: (
"gpt_neox.layers.{bid}.mlp.dense_h_to_4h", # gptneox
- "transformer.h.{bid}.mlp.c_fc", # gpt2
+ "transformer.h.{bid}.mlp.c_fc", # gpt2 jais
"transformer.blocks.{bid}.ffn.up_proj", # mpt
"transformer.h.{bid}.mlp.dense_h_to_4h", # falcon
"h.{bid}.mlp.dense_h_to_4h", # bloom
@@ -246,6 +276,7 @@ class TensorNameMap:
"model.layers.{bid}.mlp.c_fc", # starcoder2
"encoder.layer.{bid}.mlp.gated_layers_v", # jina-bert-v2
"model.layers.{bid}.residual_mlp.w3", # arctic
+ "encoder.layers.{bid}.mlp.dense_h_to_4h", # chatglm
),
MODEL_TENSOR.FFN_UP_EXP: (
@@ -270,6 +301,7 @@ class TensorNameMap:
"model.layers.{bid}.mlp.gate_proj", # llama-hf refact
"layers.{bid}.feed_forward.w1", # llama-pth
"transformer.h.{bid}.mlp.w2", # qwen
+ "transformer.h.{bid}.mlp.c_fc2", # jais
"model.layers.layers.{bid}.mlp.gate_proj", # plamo
"model.layers.{bid}.feed_forward.w1", # internlm2
"encoder.layers.{bid}.mlp.fc12", # nomic-bert
@@ -293,7 +325,7 @@ class TensorNameMap:
# Feed-forward down
MODEL_TENSOR.FFN_DOWN: (
"gpt_neox.layers.{bid}.mlp.dense_4h_to_h", # gptneox
- "transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen
+ "transformer.h.{bid}.mlp.c_proj", # gpt2 refact qwen jais
"transformer.blocks.{bid}.ffn.down_proj", # mpt
"transformer.h.{bid}.mlp.dense_4h_to_h", # falcon
"h.{bid}.mlp.dense_4h_to_h", # bloom
@@ -311,8 +343,10 @@ class TensorNameMap:
"encoder.layers.{bid}.mlp.fc2", # nomic-bert
"model.layers.{bid}.mlp.c_proj", # starcoder2
"encoder.layer.{bid}.mlp.wo", # jina-bert-v2
+ "transformer.layers.{bid}.ffn.proj_2", # openelm
"model.layers.{bid}.residual_mlp.w2", # arctic
"encoder.layer.{bid}.mlp.down_layer", # jina-bert-v2
+ "encoder.layers.{bid}.mlp.dense_4h_to_h", # chatglm
),
MODEL_TENSOR.FFN_DOWN_EXP: (
@@ -332,7 +366,8 @@ class TensorNameMap:
"model.layers.{bid}.self_attn.q_layernorm", # persimmon
"model.layers.{bid}.self_attn.q_norm", # cohere
"transformer.blocks.{bid}.attn.q_ln", # sea-lion
- "encoder.layer.{bid}.attention.self.layer_norm_q" # jina-bert-v2
+ "encoder.layer.{bid}.attention.self.layer_norm_q", # jina-bert-v2
+ "transformer.layers.{bid}.attn.q_norm", # openelm
),
MODEL_TENSOR.ATTN_K_NORM: (
@@ -340,7 +375,8 @@ class TensorNameMap:
"model.layers.{bid}.self_attn.k_layernorm", # persimmon
"model.layers.{bid}.self_attn.k_norm", # cohere
"transformer.blocks.{bid}.attn.k_ln", # sea-lion
- "encoder.layer.{bid}.attention.self.layer_norm_k" # jina-bert-v2
+ "encoder.layer.{bid}.attention.self.layer_norm_k", # jina-bert-v2
+ "transformer.layers.{bid}.attn.k_norm", # openelm
),
MODEL_TENSOR.ROPE_FREQS: (
@@ -421,6 +457,120 @@ class TensorNameMap:
MODEL_TENSOR.FFN_SUB_NORM: (
"model.layers.{bid}.mlp.ffn_layernorm", # bitnet
),
+
+ MODEL_TENSOR.DEC_ATTN_NORM: (
+ "decoder.block.{bid}.layer.0.layer_norm", # t5
+ ),
+
+ MODEL_TENSOR.DEC_ATTN_Q: (
+ "decoder.block.{bid}.layer.0.SelfAttention.q", # t5
+ ),
+
+ MODEL_TENSOR.DEC_ATTN_K: (
+ "decoder.block.{bid}.layer.0.SelfAttention.k", # t5
+ ),
+
+ MODEL_TENSOR.DEC_ATTN_V: (
+ "decoder.block.{bid}.layer.0.SelfAttention.v", # t5
+ ),
+
+ MODEL_TENSOR.DEC_ATTN_OUT: (
+ "decoder.block.{bid}.layer.0.SelfAttention.o", # t5
+ ),
+
+ MODEL_TENSOR.DEC_ATTN_REL_B: (
+ "decoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
+ ),
+
+ MODEL_TENSOR.DEC_CROSS_ATTN_NORM: (
+ "decoder.block.{bid}.layer.1.layer_norm", # t5
+ ),
+
+ MODEL_TENSOR.DEC_CROSS_ATTN_Q: (
+ "decoder.block.{bid}.layer.1.EncDecAttention.q", # t5
+ ),
+
+ MODEL_TENSOR.DEC_CROSS_ATTN_K: (
+ "decoder.block.{bid}.layer.1.EncDecAttention.k", # t5
+ ),
+
+ MODEL_TENSOR.DEC_CROSS_ATTN_V: (
+ "decoder.block.{bid}.layer.1.EncDecAttention.v", # t5
+ ),
+
+ MODEL_TENSOR.DEC_CROSS_ATTN_OUT: (
+ "decoder.block.{bid}.layer.1.EncDecAttention.o", # t5
+ ),
+
+ MODEL_TENSOR.DEC_CROSS_ATTN_REL_B: (
+ "decoder.block.{bid}.layer.1.EncDecAttention.relative_attention_bias", # t5
+ ),
+
+ MODEL_TENSOR.DEC_FFN_NORM: (
+ "decoder.block.{bid}.layer.2.layer_norm", # t5
+ ),
+
+ MODEL_TENSOR.DEC_FFN_GATE: (
+ "decoder.block.{bid}.layer.2.DenseReluDense.wi_0", # flan-t5
+ ),
+
+ MODEL_TENSOR.DEC_FFN_UP: (
+ "decoder.block.{bid}.layer.2.DenseReluDense.wi", # t5
+ "decoder.block.{bid}.layer.2.DenseReluDense.wi_1", # flan-t5
+ ),
+
+ MODEL_TENSOR.DEC_FFN_DOWN: (
+ "decoder.block.{bid}.layer.2.DenseReluDense.wo", # t5
+ ),
+
+ MODEL_TENSOR.DEC_OUTPUT_NORM: (
+ "decoder.final_layer_norm", # t5
+ ),
+
+ MODEL_TENSOR.ENC_ATTN_NORM: (
+ "encoder.block.{bid}.layer.0.layer_norm", # t5
+ ),
+
+ MODEL_TENSOR.ENC_ATTN_Q: (
+ "encoder.block.{bid}.layer.0.SelfAttention.q", # t5
+ ),
+
+ MODEL_TENSOR.ENC_ATTN_K: (
+ "encoder.block.{bid}.layer.0.SelfAttention.k", # t5
+ ),
+
+ MODEL_TENSOR.ENC_ATTN_V: (
+ "encoder.block.{bid}.layer.0.SelfAttention.v", # t5
+ ),
+
+ MODEL_TENSOR.ENC_ATTN_OUT: (
+ "encoder.block.{bid}.layer.0.SelfAttention.o", # t5
+ ),
+
+ MODEL_TENSOR.ENC_ATTN_REL_B: (
+ "encoder.block.{bid}.layer.0.SelfAttention.relative_attention_bias", # t5
+ ),
+
+ MODEL_TENSOR.ENC_FFN_NORM: (
+ "encoder.block.{bid}.layer.1.layer_norm", # t5
+ ),
+
+ MODEL_TENSOR.ENC_FFN_GATE: (
+ "encoder.block.{bid}.layer.1.DenseReluDense.wi_0", # flan-t5
+ ),
+
+ MODEL_TENSOR.ENC_FFN_UP: (
+ "encoder.block.{bid}.layer.1.DenseReluDense.wi", # t5
+ "encoder.block.{bid}.layer.1.DenseReluDense.wi_1", # flan-t5
+ ),
+
+ MODEL_TENSOR.ENC_FFN_DOWN: (
+ "encoder.block.{bid}.layer.1.DenseReluDense.wo", # t5
+ ),
+
+ MODEL_TENSOR.ENC_OUTPUT_NORM: (
+ "encoder.final_layer_norm", # t5
+ ),
}
# architecture-specific block mappings
@@ -452,14 +602,12 @@ class TensorNameMap:
for tensor, keys in self.block_mappings_cfg.items():
if tensor not in MODEL_TENSORS[arch]:
continue
- # TODO: make this configurable
- n_experts = 160
- for xid in range(n_experts):
- tensor_name = TENSOR_NAMES[tensor].format(bid = bid, xid = xid)
- self.mapping[tensor_name] = (tensor, tensor_name)
- for key in keys:
- key = key.format(bid = bid, xid = xid)
- self.mapping[key] = (tensor, tensor_name)
+
+ tensor_name = TENSOR_NAMES[tensor].format(bid = bid)
+ self.mapping[tensor_name] = (tensor, tensor_name)
+ for key in keys:
+ key = key.format(bid = bid)
+ self.mapping[key] = (tensor, tensor_name)
def get_type_and_name(self, key: str, try_suffixes: Sequence[str] = ()) -> tuple[MODEL_TENSOR, str] | None:
result = self.mapping.get(key)
diff --git a/gguf-py/gguf/utility.py b/gguf-py/gguf/utility.py
new file mode 100644
index 00000000..40d59b75
--- /dev/null
+++ b/gguf-py/gguf/utility.py
@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from typing import Literal
+
+
+def fill_templated_filename(filename: str, output_type: str | None) -> str:
+ # Given a file name fill in any type templates e.g. 'some-model-name.{ftype}.gguf'
+ ftype_lowercase: str = output_type.lower() if output_type is not None else ""
+ ftype_uppercase: str = output_type.upper() if output_type is not None else ""
+ return filename.format(ftype_lowercase,
+ outtype=ftype_lowercase, ftype=ftype_lowercase,
+ OUTTYPE=ftype_uppercase, FTYPE=ftype_uppercase)
+
+
+def model_weight_count_rounded_notation(model_params_count: int, min_digits: int = 2) -> str:
+ if model_params_count > 1e12 :
+ # Trillions Of Parameters
+ scaled_model_params = model_params_count * 1e-12
+ scale_suffix = "T"
+ elif model_params_count > 1e9 :
+ # Billions Of Parameters
+ scaled_model_params = model_params_count * 1e-9
+ scale_suffix = "B"
+ elif model_params_count > 1e6 :
+ # Millions Of Parameters
+ scaled_model_params = model_params_count * 1e-6
+ scale_suffix = "M"
+ else:
+ # Thousands Of Parameters
+ scaled_model_params = model_params_count * 1e-3
+ scale_suffix = "K"
+
+ fix = max(min_digits - len(str(round(scaled_model_params)).lstrip('0')), 0)
+
+ return f"{scaled_model_params:.{fix}f}{scale_suffix}"
+
+
+def size_label(total_params: int, shared_params: int, expert_params: int, expert_count: int) -> str:
+
+ if expert_count > 0:
+ pretty_size = model_weight_count_rounded_notation(abs(shared_params) + abs(expert_params), min_digits=2)
+ size_class = f"{expert_count}x{pretty_size}"
+ else:
+ size_class = model_weight_count_rounded_notation(abs(total_params), min_digits=2)
+
+ return size_class
+
+
+def naming_convention(model_name: str | None, base_name: str | None, finetune_string: str | None, version_string: str | None, size_label: str | None, output_type: str | None, model_type: Literal['vocab', 'LoRA'] | None = None) -> str:
+ # Reference: https://github.com/ggerganov/ggml/blob/master/docs/gguf.md#gguf-naming-convention
+
+ if base_name is not None:
+ name = base_name.strip().replace(' ', '-').replace('/', '-')
+ elif model_name is not None:
+ name = model_name.strip().replace(' ', '-').replace('/', '-')
+ else:
+ name = "ggml-model"
+
+ parameters = f"-{size_label}" if size_label is not None else ""
+
+ finetune = f"-{finetune_string.strip().replace(' ', '-')}" if finetune_string is not None else ""
+
+ version = f"-{version_string.strip().replace(' ', '-')}" if version_string is not None else ""
+
+ encoding = f"-{output_type.strip().replace(' ', '-').upper()}" if output_type is not None else ""
+
+ kind = f"-{model_type.strip().replace(' ', '-')}" if model_type is not None else ""
+
+ return f"{name}{parameters}{finetune}{version}{encoding}{kind}"
diff --git a/gguf-py/pyproject.toml b/gguf-py/pyproject.toml
index 36e63ee3..19f6761e 100644
--- a/gguf-py/pyproject.toml
+++ b/gguf-py/pyproject.toml
@@ -1,6 +1,6 @@
[tool.poetry]
name = "gguf"
-version = "0.9.0"
+version = "0.9.1"
description = "Read and write ML models in GGUF for GGML"
authors = ["GGML <ggml@ggml.ai>"]
packages = [
@@ -22,6 +22,7 @@ classifiers = [
python = ">=3.8"
numpy = ">=1.17"
tqdm = ">=4.27"
+pyyaml = ">=5.1"
[tool.poetry.dev-dependencies]
pytest = "^5.2"
diff --git a/gguf-py/scripts/__init__.py b/gguf-py/scripts/__init__.py
index 1ad45639..e77f2e9c 100644
--- a/gguf-py/scripts/__init__.py
+++ b/gguf-py/scripts/__init__.py
@@ -1,13 +1,6 @@
-import os
+# pyright: reportUnusedImport=false
-from importlib import import_module
-
-
-os.environ["NO_LOCAL_GGUF"] = "TRUE"
-
-gguf_convert_endian_entrypoint = import_module("scripts.gguf-convert-endian").main
-gguf_dump_entrypoint = import_module("scripts.gguf-dump").main
-gguf_set_metadata_entrypoint = import_module("scripts.gguf-set-metadata").main
-gguf_new_metadata_entrypoint = import_module("scripts.gguf-new-metadata").main
-
-del import_module, os
+from .gguf_convert_endian import main as gguf_convert_endian_entrypoint
+from .gguf_dump import main as gguf_dump_entrypoint
+from .gguf_set_metadata import main as gguf_set_metadata_entrypoint
+from .gguf_new_metadata import main as gguf_new_metadata_entrypoint
diff --git a/gguf-py/scripts/gguf-convert-endian.py b/gguf-py/scripts/gguf_convert_endian.py
index b698af0f..b698af0f 100755
--- a/gguf-py/scripts/gguf-convert-endian.py
+++ b/gguf-py/scripts/gguf_convert_endian.py
diff --git a/gguf-py/scripts/gguf-dump.py b/gguf-py/scripts/gguf_dump.py
index 92d14d6c..1b654654 100755
--- a/gguf-py/scripts/gguf-dump.py
+++ b/gguf-py/scripts/gguf_dump.py
@@ -4,6 +4,7 @@ from __future__ import annotations
import logging
import argparse
import os
+import re
import sys
from pathlib import Path
from typing import Any
@@ -208,7 +209,9 @@ def translate_tensor_name(name):
'ssm_d': 'State space model skip connection',
'ssm_dt': 'State space model time step',
'ssm_out': 'State space model output projection',
- 'blk': 'Block'
+ 'blk': 'Block',
+ 'enc': 'Encoder',
+ 'dec': 'Decoder',
}
expanded_words = []
@@ -242,26 +245,58 @@ def dump_markdown_metadata(reader: GGUFReader, args: argparse.Namespace) -> None
else:
pretty_type = str(field.types[-1].name)
+ def escape_markdown_inline_code(value_string):
+ # Find the longest contiguous sequence of backticks in the string then
+ # wrap string with appropriate number of backticks required to escape it
+ max_backticks = max((len(match.group(0)) for match in re.finditer(r'`+', value_string)), default=0)
+ inline_code_marker = '`' * (max_backticks + 1)
+
+ # If the string starts or ends with a backtick, add a space at the beginning and end
+ if value_string.startswith('`') or value_string.endswith('`'):
+ value_string = f" {value_string} "
+
+ return f"{inline_code_marker}{value_string}{inline_code_marker}"
+
total_elements = len(field.data)
value = ""
if len(field.types) == 1:
curr_type = field.types[0]
if curr_type == GGUFValueType.STRING:
- value = repr(str(bytes(field.parts[-1]), encoding='utf-8')[:60])
+ truncate_length = 60
+ value_string = str(bytes(field.parts[-1]), encoding='utf-8')
+ if len(value_string) > truncate_length:
+ head = escape_markdown_inline_code(value_string[:truncate_length // 2])
+ tail = escape_markdown_inline_code(value_string[-truncate_length // 2:])
+ value = "{head}...{tail}".format(head=head, tail=tail)
+ else:
+ value = escape_markdown_inline_code(value_string)
elif curr_type in reader.gguf_scalar_to_np:
value = str(field.parts[-1][0])
else:
if field.types[0] == GGUFValueType.ARRAY:
curr_type = field.types[1]
+ array_elements = []
+
if curr_type == GGUFValueType.STRING:
render_element = min(5, total_elements)
for element_pos in range(render_element):
- value += repr(str(bytes(field.parts[-1 - element_pos]), encoding='utf-8')[:5]) + (", " if total_elements > 1 else "")
+ truncate_length = 30
+ value_string = str(bytes(field.parts[-1 - (total_elements - element_pos - 1) * 2]), encoding='utf-8')
+ if len(value_string) > truncate_length:
+ head = escape_markdown_inline_code(value_string[:truncate_length // 2])
+ tail = escape_markdown_inline_code(value_string[-truncate_length // 2:])
+ value = "{head}...{tail}".format(head=head, tail=tail)
+ else:
+ value = escape_markdown_inline_code(value_string)
+ array_elements.append(value)
+
elif curr_type in reader.gguf_scalar_to_np:
render_element = min(7, total_elements)
for element_pos in range(render_element):
- value += str(field.parts[-1 - element_pos][0]) + (", " if total_elements > 1 else "")
- value = f'[ {value}{" ..." if total_elements > 1 else ""} ]'
+ array_elements.append(str(field.parts[-1 - (total_elements - element_pos - 1)][0]))
+
+ value = f'[ {", ".join(array_elements).strip()}{", ..." if total_elements > len(array_elements) else ""} ]'
+
kv_dump_table.append({"n":n, "pretty_type":pretty_type, "total_elements":total_elements, "field_name":field.name, "value":value})
kv_dump_table_header_map = [
@@ -291,6 +326,10 @@ def dump_markdown_metadata(reader: GGUFReader, args: argparse.Namespace) -> None
tensor_group_name = "base"
if tensor_components[0] == 'blk':
tensor_group_name = f"{tensor_components[0]}.{tensor_components[1]}"
+ elif tensor_components[0] in ['enc', 'dec'] and tensor_components[1] == 'blk':
+ tensor_group_name = f"{tensor_components[0]}.{tensor_components[1]}.{tensor_components[2]}"
+ elif tensor_components[0] in ['enc', 'dec']:
+ tensor_group_name = f"{tensor_components[0]}"
# Check if new Tensor Group
if tensor_group_name not in tensor_groups:
@@ -313,6 +352,27 @@ def dump_markdown_metadata(reader: GGUFReader, args: argparse.Namespace) -> None
markdown_content += "\n"
+ markdown_content += "### Tensor Data Offset\n"
+ markdown_content += '\n'
+ markdown_content += 'This table contains the offset and data segment relative to start of file\n'
+ markdown_content += '\n'
+
+ tensor_mapping_table: list[dict[str, str | int]] = []
+ for key, tensor in enumerate(reader.tensors):
+ data_offset_pretty = '{0:#16x}'.format(tensor.data_offset)
+ data_size_pretty = '{0:#16x}'.format(tensor.n_bytes)
+ tensor_mapping_table.append({"t_id":key, "layer_name":tensor.name, "data_offset":data_offset_pretty, "data_size":data_size_pretty})
+
+ tensors_mapping_table_header_map = [
+ {'key_name':'t_id', 'header_name':'T_ID', 'align':'right'},
+ {'key_name':'layer_name', 'header_name':'Tensor Layer Name', 'align':'left'},
+ {'key_name':'data_offset', 'header_name':'Data Offset (B)', 'align':'right'},
+ {'key_name':'data_size', 'header_name':'Data Size (B)', 'align':'right'},
+ ]
+
+ markdown_content += markdown_table_with_alignment_support(tensors_mapping_table_header_map, tensor_mapping_table)
+ markdown_content += "\n"
+
for group in tensor_prefix_order:
tensors = tensor_groups[group]
group_elements = sum(tensor.n_elements for tensor in tensors)
@@ -355,7 +415,7 @@ def dump_markdown_metadata(reader: GGUFReader, args: argparse.Namespace) -> None
markdown_content += f"- Percentage of total elements: {group_percentage:.2f}%\n"
markdown_content += "\n\n"
- print(markdown_content) # noqa: NP100
+ print(markdown_content) # noqa: NP100
def main() -> None:
@@ -364,6 +424,8 @@ def main() -> None:
parser.add_argument("--no-tensors", action="store_true", help="Don't dump tensor metadata")
parser.add_argument("--json", action="store_true", help="Produce JSON output")
parser.add_argument("--json-array", action="store_true", help="Include full array values in JSON output (long)")
+ parser.add_argument("--data-offset", action="store_true", help="Start of data offset")
+ parser.add_argument("--data-alignment", action="store_true", help="Data alignment applied globally to data field")
parser.add_argument("--markdown", action="store_true", help="Produce markdown output")
parser.add_argument("--verbose", action="store_true", help="increase output verbosity")
@@ -371,7 +433,7 @@ def main() -> None:
logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
- if not args.json and not args.markdown:
+ if not args.json and not args.markdown and not args.data_offset and not args.data_alignment:
logger.info(f'* Loading: {args.model}')
reader = GGUFReader(args.model, 'r')
@@ -380,6 +442,10 @@ def main() -> None:
dump_metadata_json(reader, args)
elif args.markdown:
dump_markdown_metadata(reader, args)
+ elif args.data_offset:
+ print(reader.data_offset) # noqa: NP100
+ elif args.data_alignment:
+ print(reader.alignment) # noqa: NP100
else:
dump_metadata(reader, args)
diff --git a/gguf-py/scripts/gguf_hash.py b/gguf-py/scripts/gguf_hash.py
new file mode 100755
index 00000000..ee34d09b
--- /dev/null
+++ b/gguf-py/scripts/gguf_hash.py
@@ -0,0 +1,102 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import uuid
+import hashlib
+
+import logging
+import argparse
+import os
+import sys
+from pathlib import Path
+
+from tqdm import tqdm
+
+# Necessary to load the local gguf package
+if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent / 'gguf-py').exists():
+ sys.path.insert(0, str(Path(__file__).parent.parent))
+
+from gguf import GGUFReader # noqa: E402
+
+
+logger = logging.getLogger("gguf-hash")
+
+# UUID_NAMESPACE_LLAMA_CPP = uuid.uuid5(uuid.NAMESPACE_URL, 'en.wikipedia.org/wiki/Llama.cpp')
+UUID_NAMESPACE_LLAMA_CPP = uuid.UUID('ef001206-dadc-5f6d-a15f-3359e577d4e5')
+
+
+# For more information about what field.parts and field.data represent,
+# please see the comments in the modify_gguf.py example.
+def gguf_hash(reader: GGUFReader, filename: str, disable_progress_bar: bool, no_layer: bool) -> None:
+ sha1 = hashlib.sha1()
+ sha256 = hashlib.sha256()
+ uuidv5_sha1 = hashlib.sha1()
+ uuidv5_sha1.update(UUID_NAMESPACE_LLAMA_CPP.bytes)
+
+ # Total Weight Calculation For Progress Bar
+ total_weights = 0
+ for n, tensor in enumerate(reader.tensors, 1):
+
+ # We don't need these
+ if tensor.name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")):
+ continue
+
+ # Calculate Tensor Volume
+ sum_weights_in_tensor = 1
+ for dim in tensor.shape:
+ sum_weights_in_tensor *= dim
+ total_weights += sum_weights_in_tensor
+
+ # Hash Progress Bar
+ bar = tqdm(desc="Hashing", total=total_weights, unit="weights", unit_scale=True, disable=disable_progress_bar)
+
+ # Hashing Process
+ for tensor in reader.tensors:
+
+ # We don't need these
+ if tensor.name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")):
+ continue
+
+ # Progressbar
+ sum_weights_in_tensor = 1
+ for dim in tensor.shape:
+ sum_weights_in_tensor *= dim
+ bar.update(sum_weights_in_tensor)
+
+ if not no_layer:
+
+ sha1_layer = hashlib.sha1()
+ sha1_layer.update(tensor.data.data)
+ print("sha1 {0} {1}:{2}".format(sha1_layer.hexdigest(), filename, tensor.name)) # noqa: NP100
+
+ sha256_layer = hashlib.sha256()
+ sha256_layer.update(tensor.data.data)
+ print("sha256 {0} {1}:{2}".format(sha256_layer.hexdigest(), filename, tensor.name)) # noqa: NP100
+
+ sha1.update(tensor.data.data)
+ sha256.update(tensor.data.data)
+ uuidv5_sha1.update(tensor.data.data)
+
+ # Flush Hash Progress Bar
+ bar.close()
+
+ # Display Hash Output
+ print("sha1 {0} {1}".format(sha1.hexdigest(), filename)) # noqa: NP100
+ print("sha256 {0} {1}".format(sha256.hexdigest(), filename)) # noqa: NP100
+ print("uuid {0} {1}".format(uuid.UUID(bytes=uuidv5_sha1.digest()[:16], version=5), filename)) # noqa: NP100
+
+
+def main() -> None:
+ parser = argparse.ArgumentParser(description="Dump GGUF file metadata")
+ parser.add_argument("model", type=str, help="GGUF format model filename")
+ parser.add_argument("--no-layer", action="store_true", help="exclude per layer hash")
+ parser.add_argument("--verbose", action="store_true", help="increase output verbosity")
+ parser.add_argument("--progressbar", action="store_true", help="enable progressbar")
+ args = parser.parse_args(None if len(sys.argv) > 1 else ["--help"])
+ logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)
+ reader = GGUFReader(args.model, 'r')
+ gguf_hash(reader, args.model, not args.progressbar, args.no_layer)
+
+
+if __name__ == '__main__':
+ main()
diff --git a/gguf-py/scripts/gguf-new-metadata.py b/gguf-py/scripts/gguf_new_metadata.py
index c4b90d58..fce52a8c 100755
--- a/gguf-py/scripts/gguf-new-metadata.py
+++ b/gguf-py/scripts/gguf_new_metadata.py
@@ -1,4 +1,6 @@
#!/usr/bin/env python3
+from __future__ import annotations
+
import logging
import argparse
import os
diff --git a/gguf-py/scripts/gguf-set-metadata.py b/gguf-py/scripts/gguf_set_metadata.py
index e35b651b..e35b651b 100755
--- a/gguf-py/scripts/gguf-set-metadata.py
+++ b/gguf-py/scripts/gguf_set_metadata.py
diff --git a/gguf-py/tests/__init__.py b/gguf-py/tests/__init__.py
new file mode 100644
index 00000000..d23ff9cb
--- /dev/null
+++ b/gguf-py/tests/__init__.py
@@ -0,0 +1 @@
+from .test_metadata import *
diff --git a/gguf-py/tests/test_gguf.py b/gguf-py/tests/test_gguf.py
deleted file mode 100644
index 0adeb7d5..00000000
--- a/gguf-py/tests/test_gguf.py
+++ /dev/null
@@ -1,7 +0,0 @@
-import gguf # noqa: F401
-
-# TODO: add tests
-
-
-def test_write_gguf() -> None:
- pass
diff --git a/gguf-py/tests/test_metadata.py b/gguf-py/tests/test_metadata.py
new file mode 100755
index 00000000..81a2a30a
--- /dev/null
+++ b/gguf-py/tests/test_metadata.py
@@ -0,0 +1,203 @@
+#!/usr/bin/env python3
+
+import unittest
+from pathlib import Path
+import os
+import sys
+
+# Necessary to load the local gguf package
+if "NO_LOCAL_GGUF" not in os.environ and (Path(__file__).parent.parent.parent / 'gguf-py').exists():
+ sys.path.insert(0, str(Path(__file__).parent.parent))
+
+import gguf
+
+
+class TestMetadataMethod(unittest.TestCase):
+
+ def test_id_to_title(self):
+ self.assertEqual(gguf.Metadata.id_to_title("Mixtral-8x7B-Instruct-v0.1"), "Mixtral 8x7B Instruct v0.1")
+ self.assertEqual(gguf.Metadata.id_to_title("Meta-Llama-3-8B"), "Meta Llama 3 8B")
+ self.assertEqual(gguf.Metadata.id_to_title("hermes-2-pro-llama-3-8b-DPO"), "Hermes 2 Pro Llama 3 8b DPO")
+
+ def test_get_model_id_components(self):
+ # This is the basic standard form with organization marker
+ self.assertEqual(gguf.Metadata.get_model_id_components("Mistral/Mixtral-8x7B-Instruct-v0.1"),
+ ('Mixtral-8x7B-Instruct-v0.1', "Mistral", 'Mixtral', 'Instruct', 'v0.1', '8x7B'))
+
+ # Similar to basic standard form but without organization marker
+ self.assertEqual(gguf.Metadata.get_model_id_components("Mixtral-8x7B-Instruct-v0.1"),
+ ('Mixtral-8x7B-Instruct-v0.1', None, 'Mixtral', 'Instruct', 'v0.1', '8x7B'))
+
+ # Missing version
+ self.assertEqual(gguf.Metadata.get_model_id_components("Mixtral-8x7B-Instruct"),
+ ('Mixtral-8x7B-Instruct', None, 'Mixtral', 'Instruct', None, '8x7B'))
+
+ # Missing finetune
+ self.assertEqual(gguf.Metadata.get_model_id_components("Mixtral-8x7B-v0.1"),
+ ('Mixtral-8x7B-v0.1', None, 'Mixtral', None, 'v0.1', '8x7B'))
+
+ # Base name and size label only
+ self.assertEqual(gguf.Metadata.get_model_id_components("Mixtral-8x7B"),
+ ('Mixtral-8x7B', None, 'Mixtral', None, None, '8x7B'))
+
+ # Base name and version only
+ self.assertEqual(gguf.Metadata.get_model_id_components("Mixtral-v0.1"),
+ ('Mixtral-v0.1', None, 'Mixtral', None, 'v0.1', None))
+
+ ## Edge Cases ##
+
+ # This is too ambiguous... best to err on caution and output nothing
+ self.assertEqual(gguf.Metadata.get_model_id_components("Mixtral"),
+ ('Mixtral', None, None, None, None, None))
+
+ # Basename has numbers mixed in and also size label provided. Must avoid capturing number in basename
+ self.assertEqual(gguf.Metadata.get_model_id_components("NousResearch/Meta-Llama-3-8B"),
+ ('Meta-Llama-3-8B', "NousResearch", 'Meta-Llama-3', None, None, '8B'))
+
+ # Non standard naming
+ self.assertEqual(gguf.Metadata.get_model_id_components("Qwen1.5-MoE-A2.7B-Chat"),
+ ('Qwen1.5-MoE-A2.7B-Chat', None, 'Qwen1.5-MoE', 'Chat', None, 'A2.7B'))
+
+ # Capture 'sub size labels' e.g. A14B in '57B-A14B' usually refers to activated params/weight count
+ self.assertEqual(gguf.Metadata.get_model_id_components("Qwen2-57B-A14B-Instruct"),
+ ('Qwen2-57B-A14B-Instruct', None, 'Qwen2', 'Instruct', None, '57B-A14B'))
+
+ # Check that it can handle a real model id with no version code
+ # Note that 4k in this string is non standard and microsoft were referring to context length rather than weight count
+ self.assertEqual(gguf.Metadata.get_model_id_components("microsoft/Phi-3-mini-4k-instruct", 4 * 10**9),
+ ('Phi-3-mini-4k-instruct', 'microsoft', 'Phi-3', '4k-instruct', None, 'mini'))
+
+ # There is some legitimate models with only thousands of parameters
+ self.assertEqual(gguf.Metadata.get_model_id_components("delphi-suite/stories-llama2-50k", 50 * 10**3),
+ ('stories-llama2-50k', 'delphi-suite', 'stories-llama2', None, None, '50K'))
+
+ # Non standard and not easy to disambiguate
+ self.assertEqual(gguf.Metadata.get_model_id_components("DeepSeek-Coder-V2-Lite-Instruct"),
+ ('DeepSeek-Coder-V2-Lite-Instruct', None, 'DeepSeek-Coder-V2-Lite', 'Instruct', None, None))
+
+ # This is a real model_id where they append 2DPO to refer to Direct Preference Optimization
+ self.assertEqual(gguf.Metadata.get_model_id_components("crestf411/daybreak-kunoichi-2dpo-7b"),
+ ('daybreak-kunoichi-2dpo-7b', 'crestf411', 'daybreak-kunoichi', '2dpo', None, '7B'))
+
+ # This is a real model id where the weight size has a decimal point
+ self.assertEqual(gguf.Metadata.get_model_id_components("Qwen2-0.5B-Instruct"),
+ ('Qwen2-0.5B-Instruct', None, 'Qwen2', 'Instruct', None, '0.5B'))
+
+ # Uses an underscore in the size label
+ self.assertEqual(gguf.Metadata.get_model_id_components("smallcloudai/Refact-1_6B-fim"),
+ ('Refact-1_6B-fim', 'smallcloudai', 'Refact', 'fim', None, '1.6B'))
+
+ # Uses Iter3 for the version
+ self.assertEqual(gguf.Metadata.get_model_id_components("UCLA-AGI/Gemma-2-9B-It-SPPO-Iter3"),
+ ('Gemma-2-9B-It-SPPO-Iter3', 'UCLA-AGI', 'Gemma-2', 'It-SPPO', 'Iter3', '9B'))
+
+ # Has two potential versions in the basename
+ self.assertEqual(gguf.Metadata.get_model_id_components("NousResearch/Hermes-2-Theta-Llama-3-8B"),
+ ('Hermes-2-Theta-Llama-3-8B', 'NousResearch', 'Hermes-2-Theta-Llama-3', None, None, '8B'))
+
+ # Potential version in the basename
+ self.assertEqual(gguf.Metadata.get_model_id_components("SeaLLMs/SeaLLMs-v3-7B-Chat"),
+ ('SeaLLMs-v3-7B-Chat', 'SeaLLMs', 'SeaLLMs-v3', 'Chat', None, '7B'))
+
+ # Underscore in the basename, and 1m for the context size
+ self.assertEqual(gguf.Metadata.get_model_id_components("internlm/internlm2_5-7b-chat-1m", 7 * 10**9),
+ ('internlm2_5-7b-chat-1m', 'internlm', 'internlm2_5', 'chat-1m', None, '7B'))
+
+ # Version before the finetune name
+ self.assertEqual(gguf.Metadata.get_model_id_components("pszemraj/jamba-900M-v0.13-KIx2"),
+ ('jamba-900M-v0.13-KIx2', 'pszemraj', 'jamba', 'KIx2', 'v0.13', '900M'))
+
+ # TODO: hf suffix which could be ignored but isn't
+ self.assertEqual(gguf.Metadata.get_model_id_components("state-spaces/mamba-2.8b-hf"),
+ ('mamba-2.8b-hf', 'state-spaces', 'mamba', 'hf', None, '2.8B'))
+
+ # Two sizes, don't merge them, the other is the number of tokens on which it was trained
+ self.assertEqual(gguf.Metadata.get_model_id_components("abacaj/llama-161M-100B", 161 * 10**6),
+ ('llama-161M-100B', 'abacaj', 'llama', '100b', None, '161M'))
+
+ # It's a trap, there is no size label
+ self.assertEqual(gguf.Metadata.get_model_id_components("SparseLLM/relu-100B", 1340 * 10**6),
+ ('relu-100B', 'SparseLLM', 'relu', '100b', None, None))
+
+ # Weird size notation
+ self.assertEqual(gguf.Metadata.get_model_id_components("bigscience/bloom-7b1-petals"),
+ ('bloom-7b1-petals', 'bigscience', 'bloom', 'petals', None, '7.1B'))
+
+ # Ignore full-text size labels when there are number-based ones, and deduplicate size labels
+ self.assertEqual(gguf.Metadata.get_model_id_components("MaziyarPanahi/GreenNode-mini-7B-multilingual-v1olet-Mistral-7B-Instruct-v0.1"),
+ ('GreenNode-mini-7B-multilingual-v1olet-Mistral-7B-Instruct-v0.1', 'MaziyarPanahi', 'GreenNode-mini', 'multilingual-v1olet-Mistral-Instruct', 'v0.1', '7B'))
+
+ # Instruct in a name without a size label
+ self.assertEqual(gguf.Metadata.get_model_id_components("mistralai/Mistral-Nemo-Instruct-2407"),
+ ('Mistral-Nemo-Instruct-2407', 'mistralai', 'Mistral-Nemo', 'Instruct', '2407', None))
+
+ # Non-obvious splitting relying on 'chat' keyword
+ self.assertEqual(gguf.Metadata.get_model_id_components("deepseek-ai/DeepSeek-V2-Chat-0628"),
+ ('DeepSeek-V2-Chat-0628', 'deepseek-ai', 'DeepSeek-V2', 'Chat', '0628', None))
+
+ # Multiple versions
+ self.assertEqual(gguf.Metadata.get_model_id_components("OpenGVLab/Mini-InternVL-Chat-2B-V1-5"),
+ ('Mini-InternVL-Chat-2B-V1-5', 'OpenGVLab', 'Mini-InternVL', 'Chat', 'V1-5', '2B'))
+
+ # TODO: DPO in the name
+ self.assertEqual(gguf.Metadata.get_model_id_components("jondurbin/bagel-dpo-2.8b-v0.2"),
+ ('bagel-dpo-2.8b-v0.2', 'jondurbin', 'bagel-dpo', None, 'v0.2', '2.8B'))
+
+ # DPO in name, but can't be used for the finetune to keep 'LLaMA-3' in the basename
+ self.assertEqual(gguf.Metadata.get_model_id_components("voxmenthe/SFR-Iterative-DPO-LLaMA-3-8B-R-unquantized"),
+ ('SFR-Iterative-DPO-LLaMA-3-8B-R-unquantized', 'voxmenthe', 'SFR-Iterative-DPO-LLaMA-3', 'R-unquantized', None, '8B'))
+
+ # Too ambiguous
+ # TODO: should "base" be a 'finetune' or 'size_label'?
+ # (in this case it should be a size label, but other models use it to signal that they are not finetuned)
+ self.assertEqual(gguf.Metadata.get_model_id_components("microsoft/Florence-2-base"),
+ ('Florence-2-base', 'microsoft', None, None, None, None))
+
+ ## Invalid cases ##
+
+ # Start with a dash and has dashes in rows
+ self.assertEqual(gguf.Metadata.get_model_id_components("mistralai/-Mistral--Nemo-Base-2407-"),
+ ('-Mistral--Nemo-Base-2407-', 'mistralai', 'Mistral-Nemo-Base', None, '2407', None))
+
+ ## LoRA ##
+
+ self.assertEqual(gguf.Metadata.get_model_id_components("Llama-3-Instruct-abliteration-LoRA-8B"),
+ ('Llama-3-Instruct-abliteration-LoRA-8B', None, 'Llama-3', 'Instruct-abliteration-LoRA', None, '8B'))
+
+ # Negative size --> output is a LoRA adaper --> prune "LoRA" out of the name to avoid redundancy with the suffix
+ self.assertEqual(gguf.Metadata.get_model_id_components("Llama-3-Instruct-abliteration-LoRA-8B", -1234),
+ ('Llama-3-Instruct-abliteration-LoRA-8B', None, 'Llama-3', 'Instruct-abliteration', None, '8B'))
+
+ def test_apply_metadata_heuristic_from_model_card(self):
+ model_card = {
+ 'tags': ['Llama-3', 'instruct', 'finetune', 'chatml', 'DPO', 'RLHF', 'gpt4', 'synthetic data', 'distillation', 'function calling', 'json mode', 'axolotl'],
+ 'model-index': [{'name': 'Mixtral-8x7B-Instruct-v0.1', 'results': []}],
+ 'language': ['en'],
+ 'datasets': ['teknium/OpenHermes-2.5'],
+ 'widget': [{'example_title': 'Hermes 2 Pro', 'messages': [{'role': 'system', 'content': 'You are a sentient, superintelligent artificial general intelligence, here to teach and assist me.'}, {'role': 'user', 'content': 'Write a short story about Goku discovering kirby has teamed up with Majin Buu to destroy the world.'}]}],
+ 'base_model': ["EmbeddedLLM/Mistral-7B-Merge-14-v0", "janai-hq/trinity-v1"]
+ }
+ got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card, None, None)
+ expect = gguf.Metadata()
+ expect.base_models=[{'name': 'Mistral 7B Merge 14 v0', 'organization': 'EmbeddedLLM', 'version': '14-v0', 'repo_url': 'https://huggingface.co/EmbeddedLLM/Mistral-7B-Merge-14-v0'}, {'name': 'Trinity v1', 'organization': 'Janai Hq', 'version': 'v1', 'repo_url': 'https://huggingface.co/janai-hq/trinity-v1'}]
+ expect.tags=['Llama-3', 'instruct', 'finetune', 'chatml', 'DPO', 'RLHF', 'gpt4', 'synthetic data', 'distillation', 'function calling', 'json mode', 'axolotl']
+ expect.languages=['en']
+ expect.datasets=['teknium/OpenHermes-2.5']
+
+ self.assertEqual(got, expect)
+
+ def test_apply_metadata_heuristic_from_hf_parameters(self):
+ hf_params = {"_name_or_path": "./hermes-2-pro-llama-3-8b-DPO"}
+ got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card=None, hf_params=hf_params, model_path=None)
+ expect = gguf.Metadata(name='Hermes 2 Pro Llama 3 8b DPO', finetune='DPO', basename='hermes-2-pro-llama-3', size_label='8B')
+ self.assertEqual(got, expect)
+
+ def test_apply_metadata_heuristic_from_model_dir(self):
+ model_dir_path = Path("./hermes-2-pro-llama-3-8b-DPO")
+ got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card=None, hf_params=None, model_path=model_dir_path)
+ expect = gguf.Metadata(name='Hermes 2 Pro Llama 3 8b DPO', finetune='DPO', basename='hermes-2-pro-llama-3', size_label='8B')
+ self.assertEqual(got, expect)
+
+
+if __name__ == "__main__":
+ unittest.main()