gguf-py: Refactor and add file reading support
This commit is contained in:
KerfuffleV2
parent
0a7c980b6f
commit
b8c80df741
17 changed files with 1504 additions and 1157 deletions
|
|
@ -16,7 +16,7 @@ import torch
|
||||||
from sentencepiece import SentencePieceProcessor # type: ignore[import]
|
from sentencepiece import SentencePieceProcessor # type: ignore[import]
|
||||||
|
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -17,7 +17,7 @@ import torch
|
||||||
from transformers import AutoTokenizer # type: ignore[import]
|
from transformers import AutoTokenizer # type: ignore[import]
|
||||||
|
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -17,7 +17,7 @@ import torch
|
||||||
from transformers import AutoTokenizer # type: ignore[import]
|
from transformers import AutoTokenizer # type: ignore[import]
|
||||||
|
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ import torch
|
||||||
from transformers import AutoTokenizer # type: ignore[import]
|
from transformers import AutoTokenizer # type: ignore[import]
|
||||||
|
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -12,29 +12,9 @@ import numpy as np
|
||||||
|
|
||||||
import os
|
import os
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
# Note: Does not support GGML_QKK_64
|
|
||||||
QK_K = 256
|
|
||||||
# Items here are (block size, type size)
|
|
||||||
GGML_QUANT_SIZES = {
|
|
||||||
gguf.GGMLQuantizationType.F32 : (1, 4),
|
|
||||||
gguf.GGMLQuantizationType.F16 : (1, 2),
|
|
||||||
gguf.GGMLQuantizationType.Q4_0 : (32, 2 + 16),
|
|
||||||
gguf.GGMLQuantizationType.Q4_1 : (32, 2 + 2 + 16),
|
|
||||||
gguf.GGMLQuantizationType.Q5_0 : (32, 2 + 4 + 16),
|
|
||||||
gguf.GGMLQuantizationType.Q5_1 : (32, 2 + 2 + 4 + 16),
|
|
||||||
gguf.GGMLQuantizationType.Q8_0 : (32, 2 + 32),
|
|
||||||
gguf.GGMLQuantizationType.Q8_1 : (32, 4 + 4 + 32),
|
|
||||||
gguf.GGMLQuantizationType.Q2_K : (256, 2 + 2 + QK_K // 16 + QK_K // 4),
|
|
||||||
gguf.GGMLQuantizationType.Q3_K : (256, 2 + QK_K // 4 + QK_K // 8 + 12),
|
|
||||||
gguf.GGMLQuantizationType.Q4_K : (256, 2 + 2 + QK_K // 2 + 12),
|
|
||||||
gguf.GGMLQuantizationType.Q5_K : (256, 2 + 2 + QK_K // 2 + QK_K // 8 + 12),
|
|
||||||
gguf.GGMLQuantizationType.Q6_K : (256, 2 + QK_K // 2 + QK_K // 4 + QK_K // 16),
|
|
||||||
gguf.GGMLQuantizationType.Q8_K : (256, 4 + QK_K + QK_K // 8),
|
|
||||||
}
|
|
||||||
|
|
||||||
class GGMLFormat(IntEnum):
|
class GGMLFormat(IntEnum):
|
||||||
GGML = 0
|
GGML = 0
|
||||||
GGMF = 1
|
GGMF = 1
|
||||||
|
|
@ -125,7 +105,7 @@ class Tensor:
|
||||||
(n_dims, name_len, dtype) = struct.unpack('<3I', data[offset:offset + 12])
|
(n_dims, name_len, dtype) = struct.unpack('<3I', data[offset:offset + 12])
|
||||||
assert n_dims >= 0 and n_dims <= 4, f'Invalid tensor dimensions {n_dims}'
|
assert n_dims >= 0 and n_dims <= 4, f'Invalid tensor dimensions {n_dims}'
|
||||||
assert name_len < 4096, 'Absurd tensor name length'
|
assert name_len < 4096, 'Absurd tensor name length'
|
||||||
quant = GGML_QUANT_SIZES.get(dtype)
|
quant = gguf.GGML_QUANT_SIZES.get(dtype)
|
||||||
assert quant is not None, 'Unknown tensor type'
|
assert quant is not None, 'Unknown tensor type'
|
||||||
(blksize, tysize) = quant
|
(blksize, tysize) = quant
|
||||||
offset += 12
|
offset += 12
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ import torch
|
||||||
from transformers import AutoTokenizer # type: ignore[import]
|
from transformers import AutoTokenizer # type: ignore[import]
|
||||||
|
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -6,7 +6,7 @@ import argparse
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from sentencepiece import SentencePieceProcessor
|
from sentencepiece import SentencePieceProcessor
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
def _flatten_dict(dct, tensors, prefix=None):
|
def _flatten_dict(dct, tensors, prefix=None):
|
||||||
|
|
|
||||||
|
|
@ -16,7 +16,7 @@ import torch
|
||||||
from transformers import AutoTokenizer # type: ignore[import]
|
from transformers import AutoTokenizer # type: ignore[import]
|
||||||
|
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
10
convert.py
10
convert.py
|
|
@ -3,11 +3,9 @@ from __future__ import annotations
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import concurrent.futures
|
import concurrent.futures
|
||||||
import copy
|
|
||||||
import enum
|
import enum
|
||||||
import faulthandler
|
import faulthandler
|
||||||
import functools
|
import functools
|
||||||
import io
|
|
||||||
import itertools
|
import itertools
|
||||||
import json
|
import json
|
||||||
import math
|
import math
|
||||||
|
|
@ -23,14 +21,14 @@ from abc import ABCMeta, abstractmethod
|
||||||
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
|
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Iterable, Literal, Sequence, TypeVar
|
from typing import IO, TYPE_CHECKING, Any, Callable, Iterable, Literal, TypeVar
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
from sentencepiece import SentencePieceProcessor # type: ignore[import]
|
from sentencepiece import SentencePieceProcessor # type: ignore[import]
|
||||||
|
|
||||||
import os
|
import os
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
|
|
@ -851,7 +849,7 @@ class OutputFile:
|
||||||
elif isinstance(vocab, BpeVocab):
|
elif isinstance(vocab, BpeVocab):
|
||||||
self.gguf.add_tokenizer_model("gpt2")
|
self.gguf.add_tokenizer_model("gpt2")
|
||||||
else:
|
else:
|
||||||
raise ValueError(f'Unknown vocab type: Not BpeVocab or SentencePieceVocab')
|
raise ValueError('Unknown vocab type: Not BpeVocab or SentencePieceVocab')
|
||||||
self.gguf.add_token_list(tokens)
|
self.gguf.add_token_list(tokens)
|
||||||
self.gguf.add_token_scores(scores)
|
self.gguf.add_token_scores(scores)
|
||||||
self.gguf.add_token_types(toktypes)
|
self.gguf.add_token_types(toktypes)
|
||||||
|
|
@ -905,7 +903,7 @@ class OutputFile:
|
||||||
return dt.quantize(arr)
|
return dt.quantize(arr)
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def write_all(fname_out: Path, ftype: GGMLFileType, params: Params, model: LazyModel, vocab: Vocab, svocab: gguf.SpecialVocab, concurrency: int = DEFAULT_CONCURRENCY, endianess=gguf.GGUFEndian.LITTLE) -> None:
|
def write_all(fname_out: Path, ftype: GGMLFileType, params: Params, model: LazyModel, vocab: Vocab, svocab: gguf.SpecialVocab, concurrency: int = DEFAULT_CONCURRENCY, endianess: gguf.GGUFEndian = gguf.GGUFEndian.LITTLE) -> None:
|
||||||
check_vocab_size(params, vocab)
|
check_vocab_size(params, vocab)
|
||||||
|
|
||||||
of = OutputFile(fname_out, endianess=endianess)
|
of = OutputFile(fname_out, endianess=endianess)
|
||||||
|
|
|
||||||
|
|
@ -9,7 +9,7 @@ import numpy as np
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
if 'NO_LOCAL_GGUF' not in os.environ:
|
if 'NO_LOCAL_GGUF' not in os.environ:
|
||||||
sys.path.insert(1, str(Path(__file__).parent / '..' / '..' / 'gguf-py' / 'gguf'))
|
sys.path.insert(1, str(Path(__file__).parent / '..' / '..' / 'gguf-py'))
|
||||||
import gguf
|
import gguf
|
||||||
|
|
||||||
# gguf constants
|
# gguf constants
|
||||||
|
|
|
||||||
|
|
@ -1 +1,5 @@
|
||||||
from .gguf import *
|
from .constants import *
|
||||||
|
from .gguf_writer import *
|
||||||
|
from .gguf_reader import *
|
||||||
|
from .tensor_mapping import *
|
||||||
|
from .vocab import *
|
||||||
|
|
|
||||||
469
gguf-py/gguf/constants.py
Normal file
469
gguf-py/gguf/constants.py
Normal file
|
|
@ -0,0 +1,469 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import sys
|
||||||
|
from enum import Enum, IntEnum, auto
|
||||||
|
from typing import Any, NamedTuple
|
||||||
|
|
||||||
|
#
|
||||||
|
# constants
|
||||||
|
#
|
||||||
|
|
||||||
|
GGUF_MAGIC = 0x46554747 # "GGUF"
|
||||||
|
GGUF_VERSION = 3
|
||||||
|
GGUF_DEFAULT_ALIGNMENT = 32
|
||||||
|
|
||||||
|
#
|
||||||
|
# metadata keys
|
||||||
|
#
|
||||||
|
|
||||||
|
class GeneralKeys(NamedTuple):
|
||||||
|
ARCHITECTURE = "general.architecture"
|
||||||
|
QUANTIZATION_VERSION = "general.quantization_version"
|
||||||
|
ALIGNMENT = "general.alignment"
|
||||||
|
NAME = "general.name"
|
||||||
|
AUTHOR = "general.author"
|
||||||
|
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"
|
||||||
|
|
||||||
|
class AttentionKeys(NamedTuple):
|
||||||
|
HEAD_COUNT = "{arch}.attention.head_count"
|
||||||
|
HEAD_COUNT_KV = "{arch}.attention.head_count_kv"
|
||||||
|
MAX_ALIBI_BIAS = "{arch}.attention.max_alibi_bias"
|
||||||
|
CLAMP_KQV = "{arch}.attention.clamp_kqv"
|
||||||
|
LAYERNORM_EPS = "{arch}.attention.layer_norm_epsilon"
|
||||||
|
LAYERNORM_RMS_EPS = "{arch}.attention.layer_norm_rms_epsilon"
|
||||||
|
|
||||||
|
class RopeKeys(NamedTuple):
|
||||||
|
DIMENSION_COUNT = "{arch}.rope.dimension_count"
|
||||||
|
FREQ_BASE = "{arch}.rope.freq_base"
|
||||||
|
SCALING_TYPE = "{arch}.rope.scaling.type"
|
||||||
|
SCALING_FACTOR = "{arch}.rope.scaling.factor"
|
||||||
|
SCALING_ORIG_CTX_LEN = "{arch}.rope.scaling.original_context_length"
|
||||||
|
SCALING_FINETUNED = "{arch}.rope.scaling.finetuned"
|
||||||
|
|
||||||
|
class TokenizerKeys(NamedTuple):
|
||||||
|
MODEL = "tokenizer.ggml.model"
|
||||||
|
LIST = "tokenizer.ggml.tokens"
|
||||||
|
TOKEN_TYPE = "tokenizer.ggml.token_type"
|
||||||
|
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"
|
||||||
|
HF_JSON = "tokenizer.huggingface.json"
|
||||||
|
RWKV = "tokenizer.rwkv.world"
|
||||||
|
|
||||||
|
class LLMKeys(NamedTuple):
|
||||||
|
CONTEXT_LENGTH = "{arch}.context_length"
|
||||||
|
EMBEDDING_LENGTH = "{arch}.embedding_length"
|
||||||
|
BLOCK_COUNT = "{arch}.block_count"
|
||||||
|
FEED_FORWARD_LENGTH = "{arch}.feed_forward_length"
|
||||||
|
USE_PARALLEL_RESIDUAL = "{arch}.use_parallel_residual"
|
||||||
|
TENSOR_DATA_LAYOUT = "{arch}.tensor_data_layout"
|
||||||
|
|
||||||
|
class Keys(NamedTuple):
|
||||||
|
GENERAL = GeneralKeys()
|
||||||
|
LLM = LLMKeys()
|
||||||
|
ATTENTION = AttentionKeys()
|
||||||
|
ROPE = RopeKeys()
|
||||||
|
TOKENIZER = TokenizerKeys()
|
||||||
|
|
||||||
|
KEY = Keys()
|
||||||
|
|
||||||
|
#
|
||||||
|
# recommended mapping of model tensor names for storage in gguf
|
||||||
|
#
|
||||||
|
|
||||||
|
|
||||||
|
class MODEL_ARCH(IntEnum):
|
||||||
|
LLAMA : int = auto()
|
||||||
|
FALCON : int = auto()
|
||||||
|
BAICHUAN : int = auto()
|
||||||
|
GPT2 : int = auto()
|
||||||
|
GPTJ : int = auto()
|
||||||
|
GPTNEOX : int = auto()
|
||||||
|
MPT : int = auto()
|
||||||
|
STARCODER : int = auto()
|
||||||
|
PERSIMMON : int = auto()
|
||||||
|
REFACT : int = auto()
|
||||||
|
BERT : int = auto()
|
||||||
|
BLOOM : int = auto()
|
||||||
|
|
||||||
|
|
||||||
|
class MODEL_TENSOR(IntEnum):
|
||||||
|
TOKEN_EMBD : int = auto()
|
||||||
|
TOKEN_EMBD_NORM : int = auto()
|
||||||
|
TOKEN_TYPES : int = auto()
|
||||||
|
POS_EMBD : int = auto()
|
||||||
|
OUTPUT : int = auto()
|
||||||
|
OUTPUT_NORM : int = auto()
|
||||||
|
ROPE_FREQS : int = auto()
|
||||||
|
ATTN_Q : int = auto()
|
||||||
|
ATTN_K : int = auto()
|
||||||
|
ATTN_V : int = auto()
|
||||||
|
ATTN_QKV : int = auto()
|
||||||
|
ATTN_OUT : int = auto()
|
||||||
|
ATTN_NORM : int = auto()
|
||||||
|
ATTN_NORM_2 : int = auto()
|
||||||
|
ATTN_ROT_EMBD : int = auto()
|
||||||
|
FFN_GATE : int = auto()
|
||||||
|
FFN_DOWN : int = auto()
|
||||||
|
FFN_UP : int = auto()
|
||||||
|
FFN_NORM : int = auto()
|
||||||
|
ATTN_Q_NORM : int = auto()
|
||||||
|
ATTN_K_NORM : int = auto()
|
||||||
|
|
||||||
|
|
||||||
|
MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
|
||||||
|
MODEL_ARCH.LLAMA: "llama",
|
||||||
|
MODEL_ARCH.FALCON: "falcon",
|
||||||
|
MODEL_ARCH.BAICHUAN: "baichuan",
|
||||||
|
MODEL_ARCH.GPT2: "gpt2",
|
||||||
|
MODEL_ARCH.GPTJ: "gptj",
|
||||||
|
MODEL_ARCH.GPTNEOX: "gptneox",
|
||||||
|
MODEL_ARCH.MPT: "mpt",
|
||||||
|
MODEL_ARCH.STARCODER: "starcoder",
|
||||||
|
MODEL_ARCH.PERSIMMON: "persimmon",
|
||||||
|
MODEL_ARCH.REFACT: "refact",
|
||||||
|
MODEL_ARCH.BERT: "bert",
|
||||||
|
MODEL_ARCH.BLOOM: "bloom",
|
||||||
|
}
|
||||||
|
|
||||||
|
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.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.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_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||||
|
MODEL_ARCH.LLAMA: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM,
|
||||||
|
MODEL_TENSOR.OUTPUT,
|
||||||
|
MODEL_TENSOR.ROPE_FREQS,
|
||||||
|
MODEL_TENSOR.ATTN_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_Q,
|
||||||
|
MODEL_TENSOR.ATTN_K,
|
||||||
|
MODEL_TENSOR.ATTN_V,
|
||||||
|
MODEL_TENSOR.ATTN_OUT,
|
||||||
|
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||||
|
MODEL_TENSOR.FFN_NORM,
|
||||||
|
MODEL_TENSOR.FFN_GATE,
|
||||||
|
MODEL_TENSOR.FFN_DOWN,
|
||||||
|
MODEL_TENSOR.FFN_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.GPTNEOX: [
|
||||||
|
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_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.FALCON: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM,
|
||||||
|
MODEL_TENSOR.OUTPUT,
|
||||||
|
MODEL_TENSOR.ATTN_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_NORM_2,
|
||||||
|
MODEL_TENSOR.ATTN_QKV,
|
||||||
|
MODEL_TENSOR.ATTN_OUT,
|
||||||
|
MODEL_TENSOR.FFN_DOWN,
|
||||||
|
MODEL_TENSOR.FFN_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.BAICHUAN: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM,
|
||||||
|
MODEL_TENSOR.OUTPUT,
|
||||||
|
MODEL_TENSOR.ROPE_FREQS,
|
||||||
|
MODEL_TENSOR.ATTN_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_Q,
|
||||||
|
MODEL_TENSOR.ATTN_K,
|
||||||
|
MODEL_TENSOR.ATTN_V,
|
||||||
|
MODEL_TENSOR.ATTN_OUT,
|
||||||
|
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||||
|
MODEL_TENSOR.FFN_NORM,
|
||||||
|
MODEL_TENSOR.FFN_GATE,
|
||||||
|
MODEL_TENSOR.FFN_DOWN,
|
||||||
|
MODEL_TENSOR.FFN_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.STARCODER: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.POS_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_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.BERT: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.TOKEN_TYPES,
|
||||||
|
MODEL_TENSOR.POS_EMBD,
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_Q,
|
||||||
|
MODEL_TENSOR.ATTN_K,
|
||||||
|
MODEL_TENSOR.ATTN_V,
|
||||||
|
MODEL_TENSOR.ATTN_OUT,
|
||||||
|
MODEL_TENSOR.FFN_NORM,
|
||||||
|
MODEL_TENSOR.FFN_DOWN,
|
||||||
|
MODEL_TENSOR.FFN_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.MPT: [
|
||||||
|
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_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.GPTJ: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM,
|
||||||
|
MODEL_TENSOR.OUTPUT,
|
||||||
|
MODEL_TENSOR.ATTN_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_Q,
|
||||||
|
MODEL_TENSOR.ATTN_K,
|
||||||
|
MODEL_TENSOR.ATTN_V,
|
||||||
|
MODEL_TENSOR.ATTN_OUT,
|
||||||
|
MODEL_TENSOR.FFN_DOWN,
|
||||||
|
MODEL_TENSOR.FFN_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.PERSIMMON: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.OUTPUT,
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM,
|
||||||
|
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_TENSOR.ATTN_Q_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_K_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.REFACT: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM,
|
||||||
|
MODEL_TENSOR.OUTPUT,
|
||||||
|
MODEL_TENSOR.ATTN_NORM,
|
||||||
|
MODEL_TENSOR.ATTN_Q,
|
||||||
|
MODEL_TENSOR.ATTN_K,
|
||||||
|
MODEL_TENSOR.ATTN_V,
|
||||||
|
MODEL_TENSOR.ATTN_OUT,
|
||||||
|
MODEL_TENSOR.FFN_NORM,
|
||||||
|
MODEL_TENSOR.FFN_GATE,
|
||||||
|
MODEL_TENSOR.FFN_DOWN,
|
||||||
|
MODEL_TENSOR.FFN_UP,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.BLOOM: [
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD,
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD_NORM,
|
||||||
|
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.GPT2: [
|
||||||
|
# TODO
|
||||||
|
],
|
||||||
|
# TODO
|
||||||
|
}
|
||||||
|
|
||||||
|
# tensors that will not be serialized
|
||||||
|
MODEL_TENSOR_SKIP: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
|
||||||
|
MODEL_ARCH.LLAMA: [
|
||||||
|
MODEL_TENSOR.ROPE_FREQS,
|
||||||
|
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.BAICHUAN: [
|
||||||
|
MODEL_TENSOR.ROPE_FREQS,
|
||||||
|
MODEL_TENSOR.ATTN_ROT_EMBD,
|
||||||
|
],
|
||||||
|
MODEL_ARCH.PERSIMMON: [
|
||||||
|
MODEL_TENSOR.ROPE_FREQS,
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
||||||
|
#
|
||||||
|
# types
|
||||||
|
#
|
||||||
|
|
||||||
|
class TokenType(IntEnum):
|
||||||
|
NORMAL = 1
|
||||||
|
UNKNOWN = 2
|
||||||
|
CONTROL = 3
|
||||||
|
USER_DEFINED = 4
|
||||||
|
UNUSED = 5
|
||||||
|
BYTE = 6
|
||||||
|
|
||||||
|
class RopeScalingType(Enum):
|
||||||
|
NONE = 'none'
|
||||||
|
LINEAR = 'linear'
|
||||||
|
YARN = 'yarn'
|
||||||
|
|
||||||
|
class GGMLQuantizationType(IntEnum):
|
||||||
|
F32 = 0
|
||||||
|
F16 = 1
|
||||||
|
Q4_0 = 2
|
||||||
|
Q4_1 = 3
|
||||||
|
Q5_0 = 6
|
||||||
|
Q5_1 = 7
|
||||||
|
Q8_0 = 8
|
||||||
|
Q8_1 = 9
|
||||||
|
Q2_K = 10
|
||||||
|
Q3_K = 11
|
||||||
|
Q4_K = 12
|
||||||
|
Q5_K = 13
|
||||||
|
Q6_K = 14
|
||||||
|
Q8_K = 15
|
||||||
|
|
||||||
|
class GGUFEndian(IntEnum):
|
||||||
|
LITTLE = 0
|
||||||
|
BIG = 1
|
||||||
|
|
||||||
|
|
||||||
|
class GGUFValueType(IntEnum):
|
||||||
|
UINT8 = 0
|
||||||
|
INT8 = 1
|
||||||
|
UINT16 = 2
|
||||||
|
INT16 = 3
|
||||||
|
UINT32 = 4
|
||||||
|
INT32 = 5
|
||||||
|
FLOAT32 = 6
|
||||||
|
BOOL = 7
|
||||||
|
STRING = 8
|
||||||
|
ARRAY = 9
|
||||||
|
UINT64 = 10
|
||||||
|
INT64 = 11
|
||||||
|
FLOAT64 = 12
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_type(val: Any) -> GGUFValueType:
|
||||||
|
if isinstance(val, str) or isinstance(val, bytes) or isinstance(val, bytearray):
|
||||||
|
return GGUFValueType.STRING
|
||||||
|
elif isinstance(val, list):
|
||||||
|
return GGUFValueType.ARRAY
|
||||||
|
elif isinstance(val, float):
|
||||||
|
return GGUFValueType.FLOAT32
|
||||||
|
elif isinstance(val, bool):
|
||||||
|
return GGUFValueType.BOOL
|
||||||
|
elif isinstance(val, int):
|
||||||
|
return GGUFValueType.INT32
|
||||||
|
# TODO: need help with 64-bit types in Python
|
||||||
|
else:
|
||||||
|
print("Unknown type: "+str(type(val)))
|
||||||
|
sys.exit()
|
||||||
|
|
||||||
|
# Note: Does not support GGML_QKK_64
|
||||||
|
QK_K = 256
|
||||||
|
# Items here are (block size, type size)
|
||||||
|
GGML_QUANT_SIZES = {
|
||||||
|
GGMLQuantizationType.F32 : (1, 4),
|
||||||
|
GGMLQuantizationType.F16 : (1, 2),
|
||||||
|
GGMLQuantizationType.Q4_0 : (32, 2 + 16),
|
||||||
|
GGMLQuantizationType.Q4_1 : (32, 2 + 2 + 16),
|
||||||
|
GGMLQuantizationType.Q5_0 : (32, 2 + 4 + 16),
|
||||||
|
GGMLQuantizationType.Q5_1 : (32, 2 + 2 + 4 + 16),
|
||||||
|
GGMLQuantizationType.Q8_0 : (32, 2 + 32),
|
||||||
|
GGMLQuantizationType.Q8_1 : (32, 4 + 4 + 32),
|
||||||
|
GGMLQuantizationType.Q2_K : (256, 2 + 2 + QK_K // 16 + QK_K // 4),
|
||||||
|
GGMLQuantizationType.Q3_K : (256, 2 + QK_K // 4 + QK_K // 8 + 12),
|
||||||
|
GGMLQuantizationType.Q4_K : (256, 2 + 2 + QK_K // 2 + 12),
|
||||||
|
GGMLQuantizationType.Q5_K : (256, 2 + 2 + QK_K // 2 + QK_K // 8 + 12),
|
||||||
|
GGMLQuantizationType.Q6_K : (256, 2 + QK_K // 2 + QK_K // 4 + QK_K // 16),
|
||||||
|
GGMLQuantizationType.Q8_K : (256, 4 + QK_K + QK_K // 8),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
# Aliases for backward compatibility.
|
||||||
|
|
||||||
|
# general
|
||||||
|
KEY_GENERAL_ARCHITECTURE = KEY.GENERAL.ARCHITECTURE
|
||||||
|
KEY_GENERAL_QUANTIZATION_VERSION = KEY.GENERAL.QUANTIZATION_VERSION
|
||||||
|
KEY_GENERAL_ALIGNMENT = KEY.GENERAL.ALIGNMENT
|
||||||
|
KEY_GENERAL_NAME = KEY.GENERAL.NAME
|
||||||
|
KEY_GENERAL_AUTHOR = KEY.GENERAL.AUTHOR
|
||||||
|
KEY_GENERAL_URL = KEY.GENERAL.URL
|
||||||
|
KEY_GENERAL_DESCRIPTION = KEY.GENERAL.DESCRIPTION
|
||||||
|
KEY_GENERAL_LICENSE = KEY.GENERAL.LICENSE
|
||||||
|
KEY_GENERAL_SOURCE_URL = KEY.GENERAL.SOURCE_URL
|
||||||
|
KEY_GENERAL_SOURCE_HF_REPO = KEY.GENERAL.SOURCE_HF_REPO
|
||||||
|
KEY_GENERAL_FILE_TYPE = KEY.GENERAL.FILE_TYPE
|
||||||
|
|
||||||
|
# LLM
|
||||||
|
KEY_CONTEXT_LENGTH = KEY.LLM.CONTEXT_LENGTH
|
||||||
|
KEY_EMBEDDING_LENGTH = KEY.LLM.EMBEDDING_LENGTH
|
||||||
|
KEY_BLOCK_COUNT = KEY.LLM.BLOCK_COUNT
|
||||||
|
KEY_FEED_FORWARD_LENGTH = KEY.LLM.FEED_FORWARD_LENGTH
|
||||||
|
KEY_USE_PARALLEL_RESIDUAL = KEY.LLM.USE_PARALLEL_RESIDUAL
|
||||||
|
KEY_TENSOR_DATA_LAYOUT = KEY.LLM.TENSOR_DATA_LAYOUT
|
||||||
|
|
||||||
|
# attention
|
||||||
|
KEY_ATTENTION_HEAD_COUNT = KEY.ATTENTION.HEAD_COUNT
|
||||||
|
KEY_ATTENTION_HEAD_COUNT_KV = KEY.ATTENTION.HEAD_COUNT_KV
|
||||||
|
KEY_ATTENTION_MAX_ALIBI_BIAS = KEY.ATTENTION.MAX_ALIBI_BIAS
|
||||||
|
KEY_ATTENTION_CLAMP_KQV = KEY.ATTENTION.CLAMP_KQV
|
||||||
|
KEY_ATTENTION_LAYERNORM_EPS = KEY.ATTENTION.LAYERNORM_EPS
|
||||||
|
KEY_ATTENTION_LAYERNORM_RMS_EPS = KEY.ATTENTION.LAYERNORM_RMS_EPS
|
||||||
|
|
||||||
|
# RoPE
|
||||||
|
KEY_ROPE_DIMENSION_COUNT = KEY.ROPE.DIMENSION_COUNT
|
||||||
|
KEY_ROPE_FREQ_BASE = KEY.ROPE.FREQ_BASE
|
||||||
|
KEY_ROPE_SCALING_TYPE = KEY.ROPE.SCALING_TYPE
|
||||||
|
KEY_ROPE_SCALING_FACTOR = KEY.ROPE.SCALING_FACTOR
|
||||||
|
KEY_ROPE_SCALING_ORIG_CTX_LEN = KEY.ROPE.SCALING_ORIG_CTX_LEN
|
||||||
|
KEY_ROPE_SCALING_FINETUNED = KEY.ROPE.SCALING_FINETUNED
|
||||||
|
|
||||||
|
# tokenization
|
||||||
|
KEY_TOKENIZER_MODEL = KEY.TOKENIZER.MODEL
|
||||||
|
KEY_TOKENIZER_LIST = KEY.TOKENIZER.LIST
|
||||||
|
KEY_TOKENIZER_TOKEN_TYPE = KEY.TOKENIZER.TOKEN_TYPE
|
||||||
|
KEY_TOKENIZER_SCORES = KEY.TOKENIZER.SCORES
|
||||||
|
KEY_TOKENIZER_MERGES = KEY.TOKENIZER.MERGES
|
||||||
|
KEY_TOKENIZER_BOS_ID = KEY.TOKENIZER.BOS_ID
|
||||||
|
KEY_TOKENIZER_EOS_ID = KEY.TOKENIZER.EOS_ID
|
||||||
|
KEY_TOKENIZER_UNK_ID = KEY.TOKENIZER.UNK_ID
|
||||||
|
KEY_TOKENIZER_SEP_ID = KEY.TOKENIZER.SEP_ID
|
||||||
|
KEY_TOKENIZER_PAD_ID = KEY.TOKENIZER.PAD_ID
|
||||||
|
KEY_TOKENIZER_HF_JSON = KEY.TOKENIZER.HF_JSON
|
||||||
|
KEY_TOKENIZER_RWKV = KEY.TOKENIZER.RWKV
|
||||||
1130
gguf-py/gguf/gguf.py
1130
gguf-py/gguf/gguf.py
File diff suppressed because it is too large
Load diff
270
gguf-py/gguf/gguf_reader.py
Normal file
270
gguf-py/gguf/gguf_reader.py
Normal file
|
|
@ -0,0 +1,270 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
from collections import OrderedDict
|
||||||
|
from typing import TypeVar, NamedTuple
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
import numpy.typing as npt
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
# Allow running file in package as a script.
|
||||||
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||||
|
|
||||||
|
from gguf.constants import (
|
||||||
|
GGUF_DEFAULT_ALIGNMENT,
|
||||||
|
GGUF_MAGIC,
|
||||||
|
GGUF_VERSION,
|
||||||
|
GGML_QUANT_SIZES,
|
||||||
|
GGMLQuantizationType,
|
||||||
|
GGUFValueType,
|
||||||
|
)
|
||||||
|
|
||||||
|
READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION]
|
||||||
|
|
||||||
|
|
||||||
|
class ReaderField(NamedTuple):
|
||||||
|
# Offset to start of this field.
|
||||||
|
offset: int
|
||||||
|
|
||||||
|
# Name of the field (not necessarily from file data).
|
||||||
|
name: str
|
||||||
|
|
||||||
|
# Data parts. Some types have multiple components, such as strings
|
||||||
|
# that consist of a length followed by the string data.
|
||||||
|
parts: [npt.NDArray] = []
|
||||||
|
|
||||||
|
# Indexes into parts that we can call the actual data. For example
|
||||||
|
# an array of strings will be populated with indexes to the actual
|
||||||
|
# string data.
|
||||||
|
data: [int] = [-1]
|
||||||
|
|
||||||
|
types: [GGUFValueType] = []
|
||||||
|
|
||||||
|
|
||||||
|
class ReaderTensor(NamedTuple):
|
||||||
|
name: str
|
||||||
|
tensor_type: GGMLQuantizationType
|
||||||
|
shape: npt.NDArray[np.uint32]
|
||||||
|
n_elements: int
|
||||||
|
n_bytes: int
|
||||||
|
data_offset: int
|
||||||
|
data: npt.NDArray
|
||||||
|
field: ReaderField
|
||||||
|
|
||||||
|
|
||||||
|
class GGUFReader:
|
||||||
|
byte_order: str = 'I'
|
||||||
|
fields: 'OrderedDict[str, ReaderField]' = {}
|
||||||
|
tensors: [ReaderTensor] = []
|
||||||
|
alignment: int = GGUF_DEFAULT_ALIGNMENT
|
||||||
|
|
||||||
|
_simple_value_map = {
|
||||||
|
GGUFValueType.UINT8: np.uint8,
|
||||||
|
GGUFValueType.INT8: np.int8,
|
||||||
|
GGUFValueType.UINT16: np.uint16,
|
||||||
|
GGUFValueType.INT16: np.int16,
|
||||||
|
GGUFValueType.UINT32: np.uint32,
|
||||||
|
GGUFValueType.INT32: np.int32,
|
||||||
|
GGUFValueType.FLOAT32: np.float32,
|
||||||
|
GGUFValueType.UINT64: np.uint64,
|
||||||
|
GGUFValueType.INT64: np.int64,
|
||||||
|
GGUFValueType.FLOAT64: np.float64,
|
||||||
|
GGUFValueType.BOOL: np.bool_,
|
||||||
|
}
|
||||||
|
|
||||||
|
_DT = TypeVar('T', bound = npt.DTypeLike)
|
||||||
|
def _get(self, offset: int, dtype: _DT, count: int = 1, override_order: None | str = None) -> 'npt.NDArray[_DT]':
|
||||||
|
end_offs = np.uint64(offset + np.uint64(dtype().nbytes * count))
|
||||||
|
return (self.data[np.uint64(offset):end_offs]
|
||||||
|
.view(dtype = dtype)[:count]
|
||||||
|
.newbyteorder(override_order or self.byte_order))
|
||||||
|
|
||||||
|
def _push_field(self, field: ReaderField, skip_sum: bool = False) -> int:
|
||||||
|
if field.name in self.fields:
|
||||||
|
raise KeyError(f'Duplicate {field.name} already in list at offset {field.offset}')
|
||||||
|
self.fields[field.name] = field
|
||||||
|
return 0 if skip_sum else sum(part.nbytes for part in field.parts)
|
||||||
|
|
||||||
|
def _get_str(self, offset: int) -> (npt.NDArray[np.uint64], npt.NDArray[np.uint8]):
|
||||||
|
slen = self._get(offset, np.uint64)
|
||||||
|
return (slen, self._get(offset + 8, np.uint8, slen[0]))
|
||||||
|
|
||||||
|
def _get_field_parts(self, orig_offs: int, raw_type: int) -> (int, [np.NDArray], [int], [GGUFValueType]):
|
||||||
|
offs = orig_offs
|
||||||
|
types = []
|
||||||
|
gtype = GGUFValueType(raw_type)
|
||||||
|
types.append(gtype)
|
||||||
|
# Handle strings.
|
||||||
|
if gtype == GGUFValueType.STRING:
|
||||||
|
parts = list(self._get_str(offs))
|
||||||
|
size = sum(part.nbytes for part in parts)
|
||||||
|
return (size, parts, [1], types)
|
||||||
|
# Check if it's a simple scalar type.
|
||||||
|
nptype = self._simple_value_map.get(gtype)
|
||||||
|
if nptype is not None:
|
||||||
|
val = self._get(offs, nptype)
|
||||||
|
return (val.nbytes, [val], [0], types)
|
||||||
|
# Handle arrays.
|
||||||
|
if gtype == GGUFValueType.ARRAY:
|
||||||
|
raw_itype = self._get(offs, np.uint32)
|
||||||
|
offs += raw_itype.nbytes
|
||||||
|
alen = self._get(offs, np.uint64)
|
||||||
|
offs += alen.nbytes
|
||||||
|
parts = [raw_itype, alen]
|
||||||
|
data_idxs = []
|
||||||
|
for idx in range(alen[0]):
|
||||||
|
curr_size, curr_parts, curr_idxs, curr_types = self._get_field_parts(offs, raw_itype[0])
|
||||||
|
if idx == 0:
|
||||||
|
types += curr_types
|
||||||
|
idxs_offs = len(parts)
|
||||||
|
parts += curr_parts
|
||||||
|
data_idxs += (idx + idxs_offs for idx in curr_idxs)
|
||||||
|
offs += curr_size
|
||||||
|
return (offs - orig_offs, parts, data_idxs, types)
|
||||||
|
# We can't deal with this one.
|
||||||
|
raise ValueError('Unknown/unhandled field type {gtype}')
|
||||||
|
|
||||||
|
def _get_tensor(self, orig_offs: int) -> ReaderField:
|
||||||
|
offs = np.uint64(orig_offs)
|
||||||
|
name_len, name_data = self._get_str(offs)
|
||||||
|
offs += name_len.nbytes + name_data.nbytes
|
||||||
|
n_dims = self._get(offs, np.uint32)
|
||||||
|
offs += n_dims.nbytes
|
||||||
|
dims = self._get(offs, np.uint64, n_dims[0])
|
||||||
|
offs += dims.nbytes
|
||||||
|
raw_dtype = self._get(offs, np.uint32)
|
||||||
|
offs += raw_dtype.nbytes
|
||||||
|
offset_tensor = self._get(offs, np.uint64)
|
||||||
|
offs += offset_tensor.nbytes
|
||||||
|
return ReaderField(
|
||||||
|
orig_offs,
|
||||||
|
str(name_data, encoding = 'utf-8'),
|
||||||
|
[name_len, name_data, n_dims, dims, raw_dtype, offset_tensor],
|
||||||
|
[1, 3, 4, 5],
|
||||||
|
)
|
||||||
|
|
||||||
|
def _build_fields(self, offs, count) -> int:
|
||||||
|
for _ in range(count):
|
||||||
|
orig_offs = offs
|
||||||
|
kv_klen, kv_kdata = self._get_str(offs)
|
||||||
|
offs += kv_klen.nbytes + kv_kdata.nbytes
|
||||||
|
raw_kv_type = self._get(offs, np.uint32)
|
||||||
|
offs += raw_kv_type.nbytes
|
||||||
|
parts = [kv_klen, kv_kdata, raw_kv_type]
|
||||||
|
idxs_offs = len(parts)
|
||||||
|
field_size, field_parts, field_idxs, field_types = self._get_field_parts(offs, raw_kv_type[0])
|
||||||
|
parts += field_parts
|
||||||
|
self._push_field(ReaderField(
|
||||||
|
orig_offs,
|
||||||
|
str(kv_kdata, encoding = 'utf-8'),
|
||||||
|
parts,
|
||||||
|
list(idx + idxs_offs for idx in field_idxs),
|
||||||
|
field_types,
|
||||||
|
), skip_sum = True)
|
||||||
|
offs += field_size
|
||||||
|
return offs
|
||||||
|
|
||||||
|
def _build_tensors_fields(self, offs, count) -> (int, [ReaderField]):
|
||||||
|
tensor_fields = []
|
||||||
|
for _ in range(count):
|
||||||
|
field = self._get_tensor(offs)
|
||||||
|
offs += sum(part.nbytes for part in field.parts)
|
||||||
|
tensor_fields.append(field)
|
||||||
|
return (offs, tensor_fields)
|
||||||
|
|
||||||
|
def _build_tensors(self, start_offs: int, fields: [ReaderField]) -> None:
|
||||||
|
tensors = []
|
||||||
|
for field in fields:
|
||||||
|
_name_len, name_data, _n_dims, dims, raw_dtype, offset_tensor = field.parts
|
||||||
|
ggml_type = GGMLQuantizationType(raw_dtype[0])
|
||||||
|
n_elems = np.prod(dims)
|
||||||
|
block_size, type_size = GGML_QUANT_SIZES[ggml_type]
|
||||||
|
n_bytes = np.uint64(np.uint64(n_elems) * np.uint64(type_size)) // np.uint64(block_size)
|
||||||
|
data_offs = start_offs + offset_tensor[0]
|
||||||
|
if ggml_type == GGMLQuantizationType.F32:
|
||||||
|
item_count = n_elems
|
||||||
|
item_type = np.float32
|
||||||
|
elif ggml_type == GGMLQuantizationType.F16:
|
||||||
|
item_count = n_elems
|
||||||
|
item_type = np.float16
|
||||||
|
else:
|
||||||
|
item_count = n_bytes
|
||||||
|
item_type = np.uint8
|
||||||
|
tensors.append(ReaderTensor(
|
||||||
|
name = str(name_data, encoding = 'utf-8'),
|
||||||
|
tensor_type = ggml_type,
|
||||||
|
shape = dims,
|
||||||
|
n_elements = n_elems,
|
||||||
|
n_bytes = n_bytes,
|
||||||
|
data_offset = data_offs,
|
||||||
|
data = self._get(data_offs, item_type, item_count),
|
||||||
|
field = field,
|
||||||
|
))
|
||||||
|
self.tensors = tensors
|
||||||
|
|
||||||
|
|
||||||
|
def __init__(self, path: os.PathLike[str] | str, mode: str = 'r') -> None:
|
||||||
|
self.data = np.memmap(path, mode = mode)
|
||||||
|
offs = 0
|
||||||
|
if self._get(offs, np.uint32, override_order = '<')[0] != GGUF_MAGIC:
|
||||||
|
raise ValueError('GGUF magic invalid')
|
||||||
|
offs += 4
|
||||||
|
temp = self._get(offs, np.uint32)
|
||||||
|
if temp[0] > 2000:
|
||||||
|
self.byte_order = 'S'
|
||||||
|
temp = temp.newbyteorder(self.byte_order)
|
||||||
|
version = temp[0]
|
||||||
|
if version not in READER_SUPPORTED_VERSIONS:
|
||||||
|
raise ValueError(f'Sorry, file appears to be version {version} which we cannot handle')
|
||||||
|
offs += self._push_field(ReaderField(offs, 'GGUF.version', [temp], [0], [GGUFValueType.UINT32]))
|
||||||
|
temp = self._get(offs, np.uint64, 2)
|
||||||
|
offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp[:1]], [0], [GGUFValueType.UINT64]))
|
||||||
|
offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp[1:]], [0], [GGUFValueType.UINT64]))
|
||||||
|
tensor_count, kv_count = temp
|
||||||
|
offs = self._build_fields(offs, kv_count)
|
||||||
|
offs, tensors_fields = self._build_tensors_fields(offs, tensor_count)
|
||||||
|
new_align = self.fields.get('general.alignment')
|
||||||
|
if new_align is not None:
|
||||||
|
if new_align.types != [GGUFValueType.UINT64]:
|
||||||
|
raise ValueError('Bad type for general.alignment field')
|
||||||
|
self.alignment = new_align.parts[-1]
|
||||||
|
padding = offs % self.alignment
|
||||||
|
if padding != 0:
|
||||||
|
offs += self.alignment - padding
|
||||||
|
self._build_tensors(offs, tensors_fields)
|
||||||
|
|
||||||
|
|
||||||
|
# Example usage:
|
||||||
|
if __name__ == "__main__":
|
||||||
|
if len(sys.argv) < 2:
|
||||||
|
print('gguf_reader: Error: Specify an input file', file = sys.stderr)
|
||||||
|
sys.exit(1)
|
||||||
|
print(f'* Loading: {sys.argv[1]}')
|
||||||
|
reader = GGUFReader(sys.argv[1], 'r')
|
||||||
|
print(f'\n* Dumping {len(reader.fields)} key/value pair(s)')
|
||||||
|
for n, field in enumerate(reader.fields.values(), 1):
|
||||||
|
if len(field.types) == 0:
|
||||||
|
pretty_type = 'N/A'
|
||||||
|
elif field.types[0] == GGUFValueType.ARRAY:
|
||||||
|
nest_count = len(field.types) - 1
|
||||||
|
pretty_type = '[' * nest_count + str(field.types[-1].name) + ']' * nest_count
|
||||||
|
else:
|
||||||
|
pretty_type = str(field.types[-1].name)
|
||||||
|
print(f' {n:5}: {pretty_type:10} | {len(field.data):8} | {field.name}', end = '')
|
||||||
|
if len(field.types) == 1:
|
||||||
|
curr_type = field.types[0]
|
||||||
|
if curr_type == GGUFValueType.STRING:
|
||||||
|
print(' = {0}'.format(repr(str(field.parts[-1], encoding='utf8')[:60])), end = '')
|
||||||
|
elif field.types[0] in reader._simple_value_map:
|
||||||
|
print(' = {0}'.format(field.parts[-1][0]), end = '')
|
||||||
|
print()
|
||||||
|
|
||||||
|
print(f'\n* Dumping {len(reader.tensors)} tensor(s)')
|
||||||
|
for n, tensor in enumerate(reader.tensors, 1):
|
||||||
|
|
||||||
|
prettydims = ', '.join('{0:5}'.format(d) for d in list(tensor.shape) + [1] * (4 - len(tensor.shape)))
|
||||||
|
print(f' {n:5}: {tensor.n_elements:10} | {prettydims} | {tensor.tensor_type.name:7} | {tensor.name}')
|
||||||
379
gguf-py/gguf/gguf_writer.py
Normal file
379
gguf-py/gguf/gguf_writer.py
Normal file
|
|
@ -0,0 +1,379 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
import struct
|
||||||
|
import tempfile
|
||||||
|
from io import BufferedWriter
|
||||||
|
from typing import Any, BinaryIO, Sequence
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from .constants import (
|
||||||
|
GGUF_DEFAULT_ALIGNMENT,
|
||||||
|
GGUF_MAGIC,
|
||||||
|
GGUF_VERSION,
|
||||||
|
KEY,
|
||||||
|
GGMLQuantizationType,
|
||||||
|
GGUFEndian,
|
||||||
|
GGUFValueType,
|
||||||
|
RopeScalingType,
|
||||||
|
TokenType,
|
||||||
|
)
|
||||||
|
|
||||||
|
class GGUFWriter:
|
||||||
|
fout: BufferedWriter
|
||||||
|
arch: str
|
||||||
|
offset_tensor = 0
|
||||||
|
data_alignment = GGUF_DEFAULT_ALIGNMENT
|
||||||
|
kv_data = b""
|
||||||
|
kv_data_count = 0
|
||||||
|
ti_data = b""
|
||||||
|
ti_data_count = 0
|
||||||
|
use_temp_file: bool
|
||||||
|
temp_file: tempfile.SpooledTemporaryFile[bytes] | None = None
|
||||||
|
tensors: list[tuple[np.ndarray[Any, Any], int]]
|
||||||
|
_simple_value_packing = {
|
||||||
|
GGUFValueType.UINT8: "B",
|
||||||
|
GGUFValueType.INT8: "b",
|
||||||
|
GGUFValueType.UINT16: "H",
|
||||||
|
GGUFValueType.INT16: "h",
|
||||||
|
GGUFValueType.UINT32: "I",
|
||||||
|
GGUFValueType.INT32: "i",
|
||||||
|
GGUFValueType.FLOAT32: "f",
|
||||||
|
GGUFValueType.UINT64: "Q",
|
||||||
|
GGUFValueType.INT64: "q",
|
||||||
|
GGUFValueType.FLOAT64: "d",
|
||||||
|
GGUFValueType.BOOL: "?",
|
||||||
|
}
|
||||||
|
|
||||||
|
def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes:
|
||||||
|
pack_prefix = ''
|
||||||
|
if not skip_pack_prefix:
|
||||||
|
pack_prefix = '<' if self.endianess == GGUFEndian.LITTLE else '>'
|
||||||
|
return struct.pack(f'{pack_prefix}{fmt}', value)
|
||||||
|
|
||||||
|
def _write_packed(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> None:
|
||||||
|
self.fout.write(self._pack(fmt, value, skip_pack_prefix))
|
||||||
|
|
||||||
|
def __init__(self, path: os.PathLike[str] | str, arch: str, use_temp_file: bool = True, endianess: GGUFEndian = GGUFEndian.LITTLE) -> None:
|
||||||
|
self.fout = open(path, "wb")
|
||||||
|
self.arch = arch
|
||||||
|
self.endianess = endianess
|
||||||
|
self.add_architecture()
|
||||||
|
self.use_temp_file = use_temp_file
|
||||||
|
self.tensors = []
|
||||||
|
print("gguf: This GGUF file is for {0} Endian only"
|
||||||
|
.format("Big" if self.endianess == GGUFEndian.BIG else "Little"))
|
||||||
|
|
||||||
|
def write_header_to_file(self) -> None:
|
||||||
|
self._write_packed("<I", GGUF_MAGIC, skip_pack_prefix = True)
|
||||||
|
self._write_packed("I", GGUF_VERSION)
|
||||||
|
self._write_packed("Q", self.ti_data_count)
|
||||||
|
self._write_packed("Q", self.kv_data_count)
|
||||||
|
self.flush()
|
||||||
|
# print("tensors " + str(self.ti_data_count) + " kv " + str(self.kv_data_count))
|
||||||
|
|
||||||
|
def write_kv_data_to_file(self) -> None:
|
||||||
|
self.fout.write(self.kv_data)
|
||||||
|
self.flush()
|
||||||
|
|
||||||
|
def write_ti_data_to_file(self) -> None:
|
||||||
|
self.fout.write(self.ti_data)
|
||||||
|
self.flush()
|
||||||
|
|
||||||
|
def add_key(self, key: str) -> None:
|
||||||
|
self.add_val(key, GGUFValueType.STRING, add_vtype=False)
|
||||||
|
|
||||||
|
def add_uint8(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.UINT8)
|
||||||
|
|
||||||
|
def add_int8(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.INT8)
|
||||||
|
|
||||||
|
def add_uint16(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.UINT16)
|
||||||
|
|
||||||
|
def add_int16(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.INT16)
|
||||||
|
|
||||||
|
def add_uint32(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.UINT32)
|
||||||
|
|
||||||
|
def add_int32(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.INT32)
|
||||||
|
|
||||||
|
def add_float32(self, key: str, val: float) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.FLOAT32)
|
||||||
|
|
||||||
|
def add_uint64(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.UINT64)
|
||||||
|
|
||||||
|
def add_int64(self, key: str, val: int) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.INT64)
|
||||||
|
|
||||||
|
def add_float64(self, key: str, val: float) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.FLOAT64)
|
||||||
|
|
||||||
|
def add_bool(self, key: str, val: bool) -> None:
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(val, GGUFValueType.BOOL)
|
||||||
|
|
||||||
|
def add_string(self, key: str, val: str) -> None:
|
||||||
|
if len(val) == 0:
|
||||||
|
return
|
||||||
|
self.add_key(key)
|
||||||
|
self.add_val(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(key)
|
||||||
|
self.add_val(val, GGUFValueType.ARRAY)
|
||||||
|
|
||||||
|
def add_val(self, val: Any, vtype: GGUFValueType | None = None, add_vtype: bool = True) -> None:
|
||||||
|
if vtype is None:
|
||||||
|
vtype = GGUFValueType.get_type(val)
|
||||||
|
|
||||||
|
if add_vtype:
|
||||||
|
self.kv_data += self._pack("I", vtype)
|
||||||
|
self.kv_data_count += 1
|
||||||
|
|
||||||
|
pack_fmt = self._simple_value_packing.get(vtype)
|
||||||
|
if pack_fmt is not None:
|
||||||
|
self.kv_data += self._pack(pack_fmt, val, skip_pack_prefix = vtype == GGUFValueType.BOOL)
|
||||||
|
elif vtype == GGUFValueType.STRING:
|
||||||
|
encoded_val = val.encode("utf8") if isinstance(val, str) else val
|
||||||
|
self.kv_data += self._pack("Q", len(encoded_val))
|
||||||
|
self.kv_data += encoded_val
|
||||||
|
elif vtype == GGUFValueType.ARRAY and isinstance(val, Sequence) and len(val) > 0:
|
||||||
|
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")
|
||||||
|
self.kv_data += self._pack("I", ltype)
|
||||||
|
self.kv_data += self._pack("Q", len(val))
|
||||||
|
for item in val:
|
||||||
|
self.add_val(item, add_vtype=False)
|
||||||
|
else:
|
||||||
|
raise ValueError("Invalid GGUF metadata value type or value")
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def ggml_pad(x: int, n: int) -> int:
|
||||||
|
return ((x + n - 1) // n) * n
|
||||||
|
|
||||||
|
def add_tensor_info(self, name: str, tensor_shape: Sequence[int], tensor_dtype: np.dtype[np.float16] | np.dtype[np.float32], tensor_nbytes: int, raw_dtype: GGMLQuantizationType | None = None) -> None:
|
||||||
|
if raw_dtype is None and tensor_dtype not in (np.float32, np.float16):
|
||||||
|
raise ValueError("Only F32 and F16 tensors are supported for now")
|
||||||
|
|
||||||
|
encoded_name = name.encode("utf8")
|
||||||
|
self.ti_data += self._pack("Q", len(encoded_name))
|
||||||
|
self.ti_data += encoded_name
|
||||||
|
n_dims = len(tensor_shape)
|
||||||
|
self.ti_data += self._pack("I", n_dims)
|
||||||
|
for i in range(n_dims):
|
||||||
|
self.ti_data += self._pack("Q", tensor_shape[n_dims - 1 - i])
|
||||||
|
if raw_dtype is None:
|
||||||
|
dtype = GGMLQuantizationType.F32 if tensor_dtype == np.float32 else GGMLQuantizationType.F16
|
||||||
|
else:
|
||||||
|
dtype = raw_dtype
|
||||||
|
self.ti_data += self._pack("I", dtype)
|
||||||
|
self.ti_data += self._pack("Q", self.offset_tensor)
|
||||||
|
self.offset_tensor += GGUFWriter.ggml_pad(tensor_nbytes, self.data_alignment)
|
||||||
|
self.ti_data_count += 1
|
||||||
|
|
||||||
|
def add_tensor(self, name: str, tensor: np.ndarray[Any, Any], raw_shape: Sequence[int] | None = None, raw_dtype: GGMLQuantizationType | None = None) -> None:
|
||||||
|
if self.endianess == GGUFEndian.BIG:
|
||||||
|
tensor.byteswap(inplace=True)
|
||||||
|
if self.use_temp_file and self.temp_file is None:
|
||||||
|
fp = tempfile.SpooledTemporaryFile(mode="w+b", max_size=256*1024*1024)
|
||||||
|
fp.seek(0)
|
||||||
|
self.temp_file = fp
|
||||||
|
|
||||||
|
shape: Sequence[int] = raw_shape if raw_shape is not None else tensor.shape
|
||||||
|
self.add_tensor_info(name, shape, tensor.dtype, tensor.nbytes, raw_dtype = raw_dtype)
|
||||||
|
|
||||||
|
pad = GGUFWriter.ggml_pad(tensor.nbytes, self.data_alignment) - tensor.nbytes
|
||||||
|
|
||||||
|
if self.temp_file is None:
|
||||||
|
self.tensors.append((tensor, pad))
|
||||||
|
return
|
||||||
|
|
||||||
|
tensor.tofile(self.temp_file)
|
||||||
|
|
||||||
|
if pad != 0:
|
||||||
|
self.temp_file.write(bytes([0] * pad))
|
||||||
|
|
||||||
|
def write_padding(self, fp: BinaryIO, n: int, align: int | None = None) -> None:
|
||||||
|
pad = GGUFWriter.ggml_pad(n, align if align is not None else self.data_alignment) - n
|
||||||
|
if pad != 0:
|
||||||
|
fp.write(bytes([0] * pad))
|
||||||
|
|
||||||
|
def write_tensor_data(self, tensor: np.ndarray[Any, Any]) -> None:
|
||||||
|
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)
|
||||||
|
|
||||||
|
def write_tensors_to_file(self) -> None:
|
||||||
|
self.write_ti_data_to_file()
|
||||||
|
|
||||||
|
self.write_padding(self.fout, self.fout.tell())
|
||||||
|
|
||||||
|
if self.temp_file is None:
|
||||||
|
for (currtensor, currpad) in self.tensors:
|
||||||
|
currtensor.tofile(self.fout)
|
||||||
|
if currpad != 0:
|
||||||
|
self.fout.write(bytes([0] * currpad))
|
||||||
|
return
|
||||||
|
|
||||||
|
self.temp_file.seek(0)
|
||||||
|
|
||||||
|
shutil.copyfileobj(self.temp_file, self.fout)
|
||||||
|
self.flush()
|
||||||
|
self.temp_file.close()
|
||||||
|
|
||||||
|
def flush(self) -> None:
|
||||||
|
self.fout.flush()
|
||||||
|
|
||||||
|
def close(self) -> None:
|
||||||
|
self.fout.close()
|
||||||
|
|
||||||
|
def add_architecture(self) -> None:
|
||||||
|
self.add_string(KEY.GENERAL.ARCHITECTURE, self.arch)
|
||||||
|
|
||||||
|
def add_author(self, author: str) -> None:
|
||||||
|
self.add_string(KEY.GENERAL.AUTHOR, author)
|
||||||
|
|
||||||
|
def add_tensor_data_layout(self, layout: str) -> None:
|
||||||
|
self.add_string(KEY.TENSOR_DATA_LAYOUT.format(arch=self.arch), layout)
|
||||||
|
|
||||||
|
def add_url(self, url: str) -> None:
|
||||||
|
self.add_string(KEY.GENERAL.URL, url)
|
||||||
|
|
||||||
|
def add_description(self, description: str) -> None:
|
||||||
|
self.add_string(KEY.GENERAL.DESCRIPTION, description)
|
||||||
|
|
||||||
|
def add_source_url(self, url: str) -> None:
|
||||||
|
self.add_string(KEY.GENERAL.SOURCE_URL, url)
|
||||||
|
|
||||||
|
def add_source_hf_repo(self, repo: str) -> None:
|
||||||
|
self.add_string(KEY.GENERAL.SOURCE_HF_REPO, repo)
|
||||||
|
|
||||||
|
def add_file_type(self, ftype: int) -> None:
|
||||||
|
self.add_uint32(KEY.GENERAL.FILE_TYPE, ftype)
|
||||||
|
|
||||||
|
def add_name(self, name: str) -> None:
|
||||||
|
self.add_string(KEY.GENERAL.NAME, name)
|
||||||
|
|
||||||
|
def add_quantization_version(self, quantization_version: GGMLQuantizationType) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.GENERAL.QUANTIZATION_VERSION, quantization_version)
|
||||||
|
|
||||||
|
def add_custom_alignment(self, alignment: int) -> None:
|
||||||
|
self.data_alignment = alignment
|
||||||
|
self.add_uint32(KEY.GENERAL.ALIGNMENT, alignment)
|
||||||
|
|
||||||
|
def add_context_length(self, length: int) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.LLM.CONTEXT_LENGTH.format(arch=self.arch), length)
|
||||||
|
|
||||||
|
def add_embedding_length(self, length: int) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.LLM.EMBEDDING_LENGTH.format(arch=self.arch), length)
|
||||||
|
|
||||||
|
def add_block_count(self, length: int) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.LLM.BLOCK_COUNT.format(arch=self.arch), length)
|
||||||
|
|
||||||
|
def add_feed_forward_length(self, length: int) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
|
||||||
|
|
||||||
|
def add_parallel_residual(self, use: bool) -> None:
|
||||||
|
self.add_bool(
|
||||||
|
KEY.LLM.USE_PARALLEL_RESIDUAL.format(arch=self.arch), use)
|
||||||
|
|
||||||
|
def add_head_count(self, count: int) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.ATTENTION.HEAD_COUNT.format(arch=self.arch), count)
|
||||||
|
|
||||||
|
def add_head_count_kv(self, count: int) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.ATTENTION.HEAD_COUNT_KV.format(arch=self.arch), count)
|
||||||
|
|
||||||
|
def add_max_alibi_bias(self, bias: float) -> None:
|
||||||
|
self.add_float32(
|
||||||
|
KEY.ATTENTION.MAX_ALIBI_BIAS.format(arch=self.arch), bias)
|
||||||
|
|
||||||
|
def add_clamp_kqv(self, value: float) -> None:
|
||||||
|
self.add_float32(
|
||||||
|
KEY.ATTENTION.CLAMP_KQV.format(arch=self.arch), value)
|
||||||
|
|
||||||
|
def add_layer_norm_eps(self, value: float) -> None:
|
||||||
|
self.add_float32(
|
||||||
|
KEY.ATTENTION.LAYERNORM_EPS.format(arch=self.arch), value)
|
||||||
|
|
||||||
|
def add_layer_norm_rms_eps(self, value: float) -> None:
|
||||||
|
self.add_float32(
|
||||||
|
KEY.ATTENTION.LAYERNORM_RMS_EPS.format(arch=self.arch), value)
|
||||||
|
|
||||||
|
def add_rope_dimension_count(self, count: int) -> None:
|
||||||
|
self.add_uint32(
|
||||||
|
KEY.ROPE.DIMENSION_COUNT.format(arch=self.arch), count)
|
||||||
|
|
||||||
|
def add_rope_freq_base(self, value: float) -> None:
|
||||||
|
self.add_float32(KEY.ROPE.FREQ_BASE.format(arch=self.arch), value)
|
||||||
|
|
||||||
|
def add_rope_scaling_type(self, value: RopeScalingType) -> None:
|
||||||
|
self.add_string(KEY.ROPE.SCALING_TYPE.format(arch=self.arch), value.value)
|
||||||
|
|
||||||
|
def add_rope_scaling_factor(self, value: float) -> None:
|
||||||
|
self.add_float32(KEY.ROPE.SCALING_FACTOR.format(arch=self.arch), value)
|
||||||
|
|
||||||
|
def add_rope_scaling_orig_ctx_len(self, value: int) -> None:
|
||||||
|
self.add_uint32(KEY.ROPE.SCALING_ORIG_CTX_LEN.format(arch=self.arch), value)
|
||||||
|
|
||||||
|
def add_rope_scaling_finetuned(self, value: bool) -> None:
|
||||||
|
self.add_bool(KEY.ROPE.SCALING_FINETUNED.format(arch=self.arch), value)
|
||||||
|
|
||||||
|
def add_tokenizer_model(self, model: str) -> None:
|
||||||
|
self.add_string(KEY.TOKENIZER.MODEL, model)
|
||||||
|
|
||||||
|
def add_token_list(self, tokens: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
|
||||||
|
self.add_array(KEY.TOKENIZER.LIST, tokens)
|
||||||
|
|
||||||
|
def add_token_merges(self, merges: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
|
||||||
|
self.add_array(KEY.TOKENIZER.MERGES, merges)
|
||||||
|
|
||||||
|
def add_token_types(self, types: Sequence[TokenType] | Sequence[int]) -> None:
|
||||||
|
self.add_array(KEY.TOKENIZER.TOKEN_TYPE, types)
|
||||||
|
|
||||||
|
def add_token_scores(self, scores: Sequence[float]) -> None:
|
||||||
|
self.add_array(KEY.TOKENIZER.SCORES, scores)
|
||||||
|
|
||||||
|
def add_bos_token_id(self, id: int) -> None:
|
||||||
|
self.add_uint32(KEY.TOKENIZER.BOS_ID, id)
|
||||||
|
|
||||||
|
def add_eos_token_id(self, id: int) -> None:
|
||||||
|
self.add_uint32(KEY.TOKENIZER.EOS_ID, id)
|
||||||
|
|
||||||
|
def add_unk_token_id(self, id: int) -> None:
|
||||||
|
self.add_uint32(KEY.TOKENIZER.UNK_ID, id)
|
||||||
|
|
||||||
|
def add_sep_token_id(self, id: int) -> None:
|
||||||
|
self.add_uint32(KEY.TOKENIZER.SEP_ID, id)
|
||||||
|
|
||||||
|
def add_pad_token_id(self, id: int) -> None:
|
||||||
|
self.add_uint32(KEY.TOKENIZER.PAD_ID, id)
|
||||||
254
gguf-py/gguf/tensor_mapping.py
Normal file
254
gguf-py/gguf/tensor_mapping.py
Normal file
|
|
@ -0,0 +1,254 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
from typing import Sequence
|
||||||
|
|
||||||
|
from .constants import MODEL_ARCH, MODEL_TENSOR, MODEL_TENSORS, TENSOR_NAMES
|
||||||
|
|
||||||
|
class TensorNameMap:
|
||||||
|
mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
|
||||||
|
# Token embeddings
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD: (
|
||||||
|
"gpt_neox.embed_in", # gptneox
|
||||||
|
"transformer.wte", # gpt2 gpt-j mpt refact
|
||||||
|
"transformer.word_embeddings", # falcon
|
||||||
|
"word_embeddings", # bloom
|
||||||
|
"model.embed_tokens", # llama-hf
|
||||||
|
"tok_embeddings", # llama-pth
|
||||||
|
"embeddings.word_embeddings", # bert
|
||||||
|
"language_model.embedding.word_embeddings", # persimmon
|
||||||
|
),
|
||||||
|
|
||||||
|
# Token type embeddings
|
||||||
|
MODEL_TENSOR.TOKEN_TYPES: (
|
||||||
|
"embeddings.token_type_embeddings", # bert
|
||||||
|
),
|
||||||
|
|
||||||
|
# Normalization of token embeddings
|
||||||
|
MODEL_TENSOR.TOKEN_EMBD_NORM: (
|
||||||
|
"word_embeddings_layernorm", # bloom
|
||||||
|
),
|
||||||
|
|
||||||
|
# Position embeddings
|
||||||
|
MODEL_TENSOR.POS_EMBD: (
|
||||||
|
"transformer.wpe", # gpt2
|
||||||
|
"embeddings.position_embeddings", # bert
|
||||||
|
),
|
||||||
|
|
||||||
|
# Output
|
||||||
|
MODEL_TENSOR.OUTPUT: (
|
||||||
|
"embed_out", # gptneox
|
||||||
|
"lm_head", # gpt2 mpt falcon llama-hf baichuan
|
||||||
|
"output", # llama-pth bloom
|
||||||
|
"word_embeddings_for_head", # persimmon
|
||||||
|
),
|
||||||
|
|
||||||
|
# Output norm
|
||||||
|
MODEL_TENSOR.OUTPUT_NORM: (
|
||||||
|
"gpt_neox.final_layer_norm", # gptneox
|
||||||
|
"transformer.ln_f", # gpt2 gpt-j falcon
|
||||||
|
"model.norm", # llama-hf baichuan
|
||||||
|
"norm", # llama-pth
|
||||||
|
"embeddings.LayerNorm", # bert
|
||||||
|
"transformer.norm_f", # mpt
|
||||||
|
"ln_f", # refact bloom
|
||||||
|
"language_model.encoder.final_layernorm", # persimmon
|
||||||
|
),
|
||||||
|
|
||||||
|
# Rope frequencies
|
||||||
|
MODEL_TENSOR.ROPE_FREQS: (
|
||||||
|
"rope.freqs", # llama-pth
|
||||||
|
),
|
||||||
|
}
|
||||||
|
|
||||||
|
block_mappings_cfg: dict[MODEL_TENSOR, tuple[str, ...]] = {
|
||||||
|
# Attention norm
|
||||||
|
MODEL_TENSOR.ATTN_NORM: (
|
||||||
|
"gpt_neox.layers.{bid}.input_layernorm", # gptneox
|
||||||
|
"transformer.h.{bid}.ln_1", # gpt2 gpt-j refact
|
||||||
|
"transformer.blocks.{bid}.norm_1", # mpt
|
||||||
|
"transformer.h.{bid}.input_layernorm", # falcon7b
|
||||||
|
"h.{bid}.input_layernorm", # bloom
|
||||||
|
"transformer.h.{bid}.ln_mlp", # falcon40b
|
||||||
|
"model.layers.{bid}.input_layernorm", # llama-hf
|
||||||
|
"layers.{bid}.attention_norm", # llama-pth
|
||||||
|
"encoder.layer.{bid}.attention.output.LayerNorm", # bert
|
||||||
|
"language_model.encoder.layers.{bid}.input_layernorm", # persimmon
|
||||||
|
"model.layers.{bid}.ln1", # yi
|
||||||
|
),
|
||||||
|
|
||||||
|
# Attention norm 2
|
||||||
|
MODEL_TENSOR.ATTN_NORM_2: (
|
||||||
|
"transformer.h.{bid}.ln_attn", # falcon40b
|
||||||
|
),
|
||||||
|
|
||||||
|
# Attention query-key-value
|
||||||
|
MODEL_TENSOR.ATTN_QKV: (
|
||||||
|
"gpt_neox.layers.{bid}.attention.query_key_value", # gptneox
|
||||||
|
"transformer.h.{bid}.attn.c_attn", # gpt2
|
||||||
|
"transformer.blocks.{bid}.attn.Wqkv", # mpt
|
||||||
|
"transformer.h.{bid}.self_attention.query_key_value", # falcon
|
||||||
|
"h.{bid}.self_attention.query_key_value", # bloom
|
||||||
|
"language_model.encoder.layers.{bid}.self_attention.query_key_value", # persimmon
|
||||||
|
),
|
||||||
|
|
||||||
|
# Attention query
|
||||||
|
MODEL_TENSOR.ATTN_Q: (
|
||||||
|
"model.layers.{bid}.self_attn.q_proj", # llama-hf
|
||||||
|
"layers.{bid}.attention.wq", # llama-pth
|
||||||
|
"encoder.layer.{bid}.attention.self.query", # bert
|
||||||
|
"transformer.h.{bid}.attn.q_proj", # gpt-j
|
||||||
|
),
|
||||||
|
|
||||||
|
# Attention key
|
||||||
|
MODEL_TENSOR.ATTN_K: (
|
||||||
|
"model.layers.{bid}.self_attn.k_proj", # llama-hf
|
||||||
|
"layers.{bid}.attention.wk", # llama-pth
|
||||||
|
"encoder.layer.{bid}.attention.self.key", # bert
|
||||||
|
"transformer.h.{bid}.attn.k_proj", # gpt-j
|
||||||
|
),
|
||||||
|
|
||||||
|
# Attention value
|
||||||
|
MODEL_TENSOR.ATTN_V: (
|
||||||
|
"model.layers.{bid}.self_attn.v_proj", # llama-hf
|
||||||
|
"layers.{bid}.attention.wv", # llama-pth
|
||||||
|
"encoder.layer.{bid}.attention.self.value", # bert
|
||||||
|
"transformer.h.{bid}.attn.v_proj", # gpt-j
|
||||||
|
),
|
||||||
|
|
||||||
|
# Attention output
|
||||||
|
MODEL_TENSOR.ATTN_OUT: (
|
||||||
|
"gpt_neox.layers.{bid}.attention.dense", # gptneox
|
||||||
|
"transformer.h.{bid}.attn.c_proj", # gpt2 refact
|
||||||
|
"transformer.blocks.{bid}.attn.out_proj", # mpt
|
||||||
|
"transformer.h.{bid}.self_attention.dense", # falcon
|
||||||
|
"h.{bid}.self_attention.dense", # bloom
|
||||||
|
"model.layers.{bid}.self_attn.o_proj", # llama-hf
|
||||||
|
"layers.{bid}.attention.wo", # llama-pth
|
||||||
|
"encoder.layer.{bid}.attention.output.dense", # bert
|
||||||
|
"transformer.h.{bid}.attn.out_proj", # gpt-j
|
||||||
|
"language_model.encoder.layers.{bid}.self_attention.dense" # persimmon
|
||||||
|
),
|
||||||
|
|
||||||
|
# Rotary embeddings
|
||||||
|
MODEL_TENSOR.ATTN_ROT_EMBD: (
|
||||||
|
"model.layers.{bid}.self_attn.rotary_emb.inv_freq", # llama-hf
|
||||||
|
"layers.{bid}.attention.inner_attention.rope.freqs", # llama-pth
|
||||||
|
),
|
||||||
|
|
||||||
|
# Feed-forward norm
|
||||||
|
MODEL_TENSOR.FFN_NORM: (
|
||||||
|
"gpt_neox.layers.{bid}.post_attention_layernorm", # gptneox
|
||||||
|
"transformer.h.{bid}.ln_2", # gpt2 refact
|
||||||
|
"h.{bid}.post_attention_layernorm", # bloom
|
||||||
|
"transformer.blocks.{bid}.norm_2", # mpt
|
||||||
|
"model.layers.{bid}.post_attention_layernorm", # llama-hf
|
||||||
|
"layers.{bid}.ffn_norm", # llama-pth
|
||||||
|
"encoder.layer.{bid}.output.LayerNorm", # bert
|
||||||
|
"language_model.encoder.layers.{bid}.post_attention_layernorm", # persimmon
|
||||||
|
"model.layers.{bid}.ln2", # yi
|
||||||
|
),
|
||||||
|
|
||||||
|
# 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.blocks.{bid}.ffn.up_proj", # mpt
|
||||||
|
"transformer.h.{bid}.mlp.dense_h_to_4h", # falcon
|
||||||
|
"h.{bid}.mlp.dense_h_to_4h", # bloom
|
||||||
|
"model.layers.{bid}.mlp.up_proj", # llama-hf refact
|
||||||
|
"layers.{bid}.feed_forward.w3", # llama-pth
|
||||||
|
"encoder.layer.{bid}.intermediate.dense", # bert
|
||||||
|
"transformer.h.{bid}.mlp.fc_in", # gpt-j
|
||||||
|
"language_model.encoder.layers.{bid}.mlp.dense_h_to_4h", # persimmon
|
||||||
|
),
|
||||||
|
|
||||||
|
# Feed-forward gate
|
||||||
|
MODEL_TENSOR.FFN_GATE: (
|
||||||
|
"model.layers.{bid}.mlp.gate_proj", # llama-hf refact
|
||||||
|
"layers.{bid}.feed_forward.w1", # llama-pth
|
||||||
|
),
|
||||||
|
|
||||||
|
# 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
|
||||||
|
"transformer.blocks.{bid}.ffn.down_proj", # mpt
|
||||||
|
"transformer.h.{bid}.mlp.dense_4h_to_h", # falcon
|
||||||
|
"h.{bid}.mlp.dense_4h_to_h", # bloom
|
||||||
|
"model.layers.{bid}.mlp.down_proj", # llama-hf
|
||||||
|
"layers.{bid}.feed_forward.w2", # llama-pth
|
||||||
|
"encoder.layer.{bid}.output.dense", # bert
|
||||||
|
"transformer.h.{bid}.mlp.fc_out", # gpt-j
|
||||||
|
"language_model.encoder.layers.{bid}.mlp.dense_4h_to_h", # persimmon
|
||||||
|
),
|
||||||
|
|
||||||
|
MODEL_TENSOR.ATTN_Q_NORM: (
|
||||||
|
"language_model.encoder.layers.{bid}.self_attention.q_layernorm",
|
||||||
|
),
|
||||||
|
|
||||||
|
MODEL_TENSOR.ATTN_K_NORM: (
|
||||||
|
"language_model.encoder.layers.{bid}.self_attention.k_layernorm",
|
||||||
|
),
|
||||||
|
|
||||||
|
MODEL_TENSOR.ROPE_FREQS: (
|
||||||
|
"language_model.encoder.layers.{bid}.self_attention.rotary_emb.inv_freq", # persimmon
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
mapping: dict[str, tuple[MODEL_TENSOR, str]]
|
||||||
|
|
||||||
|
def __init__(self, arch: MODEL_ARCH, n_blocks: int):
|
||||||
|
self.mapping = {}
|
||||||
|
for tensor, keys in self.mappings_cfg.items():
|
||||||
|
if tensor not in MODEL_TENSORS[arch]:
|
||||||
|
continue
|
||||||
|
tensor_name = TENSOR_NAMES[tensor]
|
||||||
|
self.mapping[tensor_name] = (tensor, tensor_name)
|
||||||
|
for key in keys:
|
||||||
|
self.mapping[key] = (tensor, tensor_name)
|
||||||
|
for bid in range(n_blocks):
|
||||||
|
for tensor, keys in self.block_mappings_cfg.items():
|
||||||
|
if tensor not in MODEL_TENSORS[arch]:
|
||||||
|
continue
|
||||||
|
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)
|
||||||
|
if result is not None:
|
||||||
|
return result
|
||||||
|
for suffix in try_suffixes:
|
||||||
|
if key.endswith(suffix):
|
||||||
|
result = self.mapping.get(key[:-len(suffix)])
|
||||||
|
if result is not None:
|
||||||
|
return (result[0], result[1] + suffix)
|
||||||
|
return None
|
||||||
|
|
||||||
|
def get_name(self, key: str, try_suffixes: Sequence[str] = ()) -> str | None:
|
||||||
|
result = self.get_type_and_name(key, try_suffixes = try_suffixes)
|
||||||
|
if result is None:
|
||||||
|
return None
|
||||||
|
return result[1]
|
||||||
|
|
||||||
|
def get_type(self, key: str, try_suffixes: Sequence[str] = ()) -> MODEL_TENSOR | None:
|
||||||
|
result = self.get_type_and_name(key, try_suffixes = try_suffixes)
|
||||||
|
if result is None:
|
||||||
|
return None
|
||||||
|
return result[0]
|
||||||
|
|
||||||
|
def __getitem__(self, key: str) -> str:
|
||||||
|
try:
|
||||||
|
return self.mapping[key][1]
|
||||||
|
except KeyError:
|
||||||
|
raise KeyError(key)
|
||||||
|
|
||||||
|
def __contains__(self, key: str) -> bool:
|
||||||
|
return key in self.mapping
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
return repr(self.mapping)
|
||||||
|
|
||||||
|
def get_tensor_name_map(arch: MODEL_ARCH, n_blocks: int) -> TensorNameMap:
|
||||||
|
return TensorNameMap(arch, n_blocks)
|
||||||
103
gguf-py/gguf/vocab.py
Normal file
103
gguf-py/gguf/vocab.py
Normal file
|
|
@ -0,0 +1,103 @@
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import sys
|
||||||
|
from pathlib import Path
|
||||||
|
from typing import Any, Callable
|
||||||
|
|
||||||
|
from .gguf_writer import GGUFWriter
|
||||||
|
|
||||||
|
class SpecialVocab:
|
||||||
|
load_merges: bool = False
|
||||||
|
merges: list[str] = []
|
||||||
|
special_token_types: tuple[str, ...] = ('bos', 'eos', 'unk', 'sep', 'pad')
|
||||||
|
special_token_ids: dict[str, int] = {}
|
||||||
|
n_vocab: int | None = None
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self, path: str | os.PathLike[str], load_merges: bool = False,
|
||||||
|
special_token_types: tuple[str, ...] | None = None,
|
||||||
|
n_vocab: int | None = None,
|
||||||
|
):
|
||||||
|
self.special_token_ids = {}
|
||||||
|
self.n_vocab = n_vocab
|
||||||
|
self.load_merges = load_merges
|
||||||
|
if special_token_types is not None:
|
||||||
|
self.special_token_types = special_token_types
|
||||||
|
self._load(Path(path))
|
||||||
|
|
||||||
|
def _load(self, path: Path) -> None:
|
||||||
|
if not self._try_load_from_tokenizer_json(path):
|
||||||
|
self._try_load_from_config_json(path)
|
||||||
|
|
||||||
|
def _set_special_token(self, typ: str, tid: Any) -> None:
|
||||||
|
if not isinstance(tid, int) or tid < 0:
|
||||||
|
return
|
||||||
|
if self.n_vocab is None or tid < self.n_vocab:
|
||||||
|
self.special_token_ids[typ] = tid
|
||||||
|
return
|
||||||
|
print(f'gguf: WARNING: Special token type {typ}, id {tid} out of range, must be under {self.n_vocab} - skipping',
|
||||||
|
file = sys.stderr)
|
||||||
|
|
||||||
|
|
||||||
|
def _try_load_from_tokenizer_json(self, path: Path) -> bool:
|
||||||
|
tokenizer_file = path / 'tokenizer.json'
|
||||||
|
if not tokenizer_file.is_file():
|
||||||
|
return False
|
||||||
|
with open(tokenizer_file, encoding = 'utf-8') as f:
|
||||||
|
tokenizer = json.load(f)
|
||||||
|
if self.load_merges:
|
||||||
|
merges = tokenizer.get('model', {}).get('merges')
|
||||||
|
if isinstance(merges, list) and len(merges) > 0 and isinstance(merges[0], str):
|
||||||
|
self.merges = merges
|
||||||
|
tokenizer_config_file = path / 'tokenizer_config.json'
|
||||||
|
added_tokens = tokenizer.get('added_tokens')
|
||||||
|
if added_tokens is None or not tokenizer_config_file.is_file():
|
||||||
|
return True
|
||||||
|
with open(tokenizer_config_file, encoding = 'utf-8') as f:
|
||||||
|
tokenizer_config = json.load(f)
|
||||||
|
for typ in self.special_token_types:
|
||||||
|
entry = tokenizer_config.get(f'{typ}_token')
|
||||||
|
if isinstance(entry, str):
|
||||||
|
tc_content = entry
|
||||||
|
elif isinstance(entry, dict):
|
||||||
|
entry_content = entry.get('content')
|
||||||
|
if not isinstance(entry_content, str):
|
||||||
|
continue
|
||||||
|
tc_content = entry_content
|
||||||
|
else:
|
||||||
|
continue
|
||||||
|
# We only need the first match here.
|
||||||
|
maybe_token_id = next((
|
||||||
|
atok.get('id') for atok in added_tokens
|
||||||
|
if atok.get('content') == tc_content), None)
|
||||||
|
self._set_special_token(typ, maybe_token_id)
|
||||||
|
return True
|
||||||
|
|
||||||
|
def _try_load_from_config_json(self, path: Path) -> bool:
|
||||||
|
config_file = path / 'config.json'
|
||||||
|
if not config_file.is_file():
|
||||||
|
return False
|
||||||
|
with open(config_file, encoding = 'utf-8') as f:
|
||||||
|
config = json.load(f)
|
||||||
|
for typ in self.special_token_types:
|
||||||
|
self._set_special_token(typ, config.get(f'{typ}_token_id'))
|
||||||
|
return True
|
||||||
|
|
||||||
|
def add_to_gguf(self, gw: GGUFWriter, quiet: bool = False) -> None:
|
||||||
|
if len(self.merges) > 0:
|
||||||
|
if not quiet:
|
||||||
|
print(f'gguf: Adding {len(self.merges)} merge(s).')
|
||||||
|
gw.add_token_merges(self.merges)
|
||||||
|
for typ, tokid in self.special_token_ids.items():
|
||||||
|
handler: Callable[[int], None] | None = getattr(gw, f'add_{typ}_token_id', None)
|
||||||
|
if handler is None:
|
||||||
|
print(f'gguf: WARNING: No handler for special token type {typ} with id {tokid} - skipping', file = sys.stderr)
|
||||||
|
continue
|
||||||
|
if not quiet:
|
||||||
|
print(f'gguf: Setting special token type {typ} to {tokid}')
|
||||||
|
handler(tokid)
|
||||||
|
|
||||||
|
def __repr__(self) -> str:
|
||||||
|
return f'<SpecialVocab with {len(self.merges)} merges and special tokens {self.special_token_ids or "unset"}>'
|
||||||
Loading…
Add table
Add a link
Reference in a new issue