fix NamedTuple and Enum usage

gguf-py/gguf/constants.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
 import sys
-from enum import Enum, IntEnum, StrEnum, auto
-from typing import Any, NamedTuple, Type
+from enum import IntEnum, StrEnum, auto
+from typing import Any, Type
 
 #
 # constants
@@ -73,16 +73,14 @@ class LLMKeys(StrEnum):
     TENSOR_DATA_LAYOUT:    str = "{arch}.tensor_data_layout"
 
 
-class Keys(NamedTuple):
-    GENERAL:   Type[GeneralKeys]   = GeneralKeys
-    LLM:       Type[LLMKeys]       = LLMKeys
-    ATTENTION: Type[AttentionKeys] = AttentionKeys
-    ROPE:      Type[RopeKeys]      = RopeKeys
-    TOKENIZER: Type[TokenizerKeys] = TokenizerKeys
+class Keys:
+    GENERAL   = GeneralKeys
+    LLM       = LLMKeys
+    ATTENTION = AttentionKeys
+    ROPE      = RopeKeys
+    TOKENIZER = TokenizerKeys
 
 
-KEY = Keys()
-
 #
 # recommended mapping of model tensor names for storage in gguf
 #
@@ -345,7 +343,7 @@ class TokenType(IntEnum):
     BYTE         = 6
 
 
-class RopeScalingType(Enum):
+class RopeScalingType(StrEnum):
     NONE   = 'none'
     LINEAR = 'linear'
     YARN   = 'yarn'
@@ -430,52 +428,52 @@ def get_type(val: Any) -> GGUFValueType:
 # Aliases for backward compatibility.
 
 # general
-KEY_GENERAL_ARCHITECTURE:         str = KEY.GENERAL.ARCHITECTURE
-KEY_GENERAL_QUANTIZATION_VERSION: str = KEY.GENERAL.QUANTIZATION_VERSION
-KEY_GENERAL_ALIGNMENT:            str = KEY.GENERAL.ALIGNMENT
-KEY_GENERAL_NAME:                 str = KEY.GENERAL.NAME
-KEY_GENERAL_AUTHOR:               str = KEY.GENERAL.AUTHOR
-KEY_GENERAL_URL:                  str = KEY.GENERAL.URL
-KEY_GENERAL_DESCRIPTION:          str = KEY.GENERAL.DESCRIPTION
-KEY_GENERAL_LICENSE:              str = KEY.GENERAL.LICENSE
-KEY_GENERAL_SOURCE_URL:           str = KEY.GENERAL.SOURCE_URL
-KEY_GENERAL_SOURCE_HF_REPO:       str = KEY.GENERAL.SOURCE_HF_REPO
-KEY_GENERAL_FILE_TYPE:            str = KEY.GENERAL.FILE_TYPE
+KEY_GENERAL_ARCHITECTURE:         str = Keys.GENERAL.ARCHITECTURE
+KEY_GENERAL_QUANTIZATION_VERSION: str = Keys.GENERAL.QUANTIZATION_VERSION
+KEY_GENERAL_ALIGNMENT:            str = Keys.GENERAL.ALIGNMENT
+KEY_GENERAL_NAME:                 str = Keys.GENERAL.NAME
+KEY_GENERAL_AUTHOR:               str = Keys.GENERAL.AUTHOR
+KEY_GENERAL_URL:                  str = Keys.GENERAL.URL
+KEY_GENERAL_DESCRIPTION:          str = Keys.GENERAL.DESCRIPTION
+KEY_GENERAL_LICENSE:              str = Keys.GENERAL.LICENSE
+KEY_GENERAL_SOURCE_URL:           str = Keys.GENERAL.SOURCE_URL
+KEY_GENERAL_SOURCE_HF_REPO:       str = Keys.GENERAL.SOURCE_HF_REPO
+KEY_GENERAL_FILE_TYPE:            str = Keys.GENERAL.FILE_TYPE
 
 # LLM
-KEY_CONTEXT_LENGTH:        str = KEY.LLM.CONTEXT_LENGTH
-KEY_EMBEDDING_LENGTH:      str = KEY.LLM.EMBEDDING_LENGTH
-KEY_BLOCK_COUNT:           str = KEY.LLM.BLOCK_COUNT
-KEY_FEED_FORWARD_LENGTH:   str = KEY.LLM.FEED_FORWARD_LENGTH
-KEY_USE_PARALLEL_RESIDUAL: str = KEY.LLM.USE_PARALLEL_RESIDUAL
-KEY_TENSOR_DATA_LAYOUT:    str = KEY.LLM.TENSOR_DATA_LAYOUT
+KEY_CONTEXT_LENGTH:        str = Keys.LLM.CONTEXT_LENGTH
+KEY_EMBEDDING_LENGTH:      str = Keys.LLM.EMBEDDING_LENGTH
+KEY_BLOCK_COUNT:           str = Keys.LLM.BLOCK_COUNT
+KEY_FEED_FORWARD_LENGTH:   str = Keys.LLM.FEED_FORWARD_LENGTH
+KEY_USE_PARALLEL_RESIDUAL: str = Keys.LLM.USE_PARALLEL_RESIDUAL
+KEY_TENSOR_DATA_LAYOUT:    str = Keys.LLM.TENSOR_DATA_LAYOUT
 
 # attention
-KEY_ATTENTION_HEAD_COUNT:        str = KEY.ATTENTION.HEAD_COUNT
-KEY_ATTENTION_HEAD_COUNT_KV:     str = KEY.ATTENTION.HEAD_COUNT_KV
-KEY_ATTENTION_MAX_ALIBI_BIAS:    str = KEY.ATTENTION.MAX_ALIBI_BIAS
-KEY_ATTENTION_CLAMP_KQV:         str = KEY.ATTENTION.CLAMP_KQV
-KEY_ATTENTION_LAYERNORM_EPS:     str = KEY.ATTENTION.LAYERNORM_EPS
-KEY_ATTENTION_LAYERNORM_RMS_EPS: str = KEY.ATTENTION.LAYERNORM_RMS_EPS
+KEY_ATTENTION_HEAD_COUNT:        str = Keys.ATTENTION.HEAD_COUNT
+KEY_ATTENTION_HEAD_COUNT_KV:     str = Keys.ATTENTION.HEAD_COUNT_KV
+KEY_ATTENTION_MAX_ALIBI_BIAS:    str = Keys.ATTENTION.MAX_ALIBI_BIAS
+KEY_ATTENTION_CLAMP_KQV:         str = Keys.ATTENTION.CLAMP_KQV
+KEY_ATTENTION_LAYERNORM_EPS:     str = Keys.ATTENTION.LAYERNORM_EPS
+KEY_ATTENTION_LAYERNORM_RMS_EPS: str = Keys.ATTENTION.LAYERNORM_RMS_EPS
 
 # RoPE
-KEY_ROPE_DIMENSION_COUNT:      str = KEY.ROPE.DIMENSION_COUNT
-KEY_ROPE_FREQ_BASE:            str = KEY.ROPE.FREQ_BASE
-KEY_ROPE_SCALING_TYPE:         str = KEY.ROPE.SCALING_TYPE
-KEY_ROPE_SCALING_FACTOR:       str = KEY.ROPE.SCALING_FACTOR
-KEY_ROPE_SCALING_ORIG_CTX_LEN: str = KEY.ROPE.SCALING_ORIG_CTX_LEN
-KEY_ROPE_SCALING_FINETUNED:    str = KEY.ROPE.SCALING_FINETUNED
+KEY_ROPE_DIMENSION_COUNT:      str = Keys.ROPE.DIMENSION_COUNT
+KEY_ROPE_FREQ_BASE:            str = Keys.ROPE.FREQ_BASE
+KEY_ROPE_SCALING_TYPE:         str = Keys.ROPE.SCALING_TYPE
+KEY_ROPE_SCALING_FACTOR:       str = Keys.ROPE.SCALING_FACTOR
+KEY_ROPE_SCALING_ORIG_CTX_LEN: str = Keys.ROPE.SCALING_ORIG_CTX_LEN
+KEY_ROPE_SCALING_FINETUNED:    str = Keys.ROPE.SCALING_FINETUNED
 
 # tokenization
-KEY_TOKENIZER_MODEL:      str = KEY.TOKENIZER.MODEL
-KEY_TOKENIZER_LIST:       str = KEY.TOKENIZER.LIST
-KEY_TOKENIZER_TOKEN_TYPE: str = KEY.TOKENIZER.TOKEN_TYPE
-KEY_TOKENIZER_SCORES:     str = KEY.TOKENIZER.SCORES
-KEY_TOKENIZER_MERGES:     str = KEY.TOKENIZER.MERGES
-KEY_TOKENIZER_BOS_ID:     str = KEY.TOKENIZER.BOS_ID
-KEY_TOKENIZER_EOS_ID:     str = KEY.TOKENIZER.EOS_ID
-KEY_TOKENIZER_UNK_ID:     str = KEY.TOKENIZER.UNK_ID
-KEY_TOKENIZER_SEP_ID:     str = KEY.TOKENIZER.SEP_ID
-KEY_TOKENIZER_PAD_ID:     str = KEY.TOKENIZER.PAD_ID
-KEY_TOKENIZER_HF_JSON:    str = KEY.TOKENIZER.HF_JSON
-KEY_TOKENIZER_RWKV:       str = KEY.TOKENIZER.RWKV
+KEY_TOKENIZER_MODEL:      str = Keys.TOKENIZER.MODEL
+KEY_TOKENIZER_LIST:       str = Keys.TOKENIZER.LIST
+KEY_TOKENIZER_TOKEN_TYPE: str = Keys.TOKENIZER.TOKEN_TYPE
+KEY_TOKENIZER_SCORES:     str = Keys.TOKENIZER.SCORES
+KEY_TOKENIZER_MERGES:     str = Keys.TOKENIZER.MERGES
+KEY_TOKENIZER_BOS_ID:     str = Keys.TOKENIZER.BOS_ID
+KEY_TOKENIZER_EOS_ID:     str = Keys.TOKENIZER.EOS_ID
+KEY_TOKENIZER_UNK_ID:     str = Keys.TOKENIZER.UNK_ID
+KEY_TOKENIZER_SEP_ID:     str = Keys.TOKENIZER.SEP_ID
+KEY_TOKENIZER_PAD_ID:     str = Keys.TOKENIZER.PAD_ID
+KEY_TOKENIZER_HF_JSON:    str = Keys.TOKENIZER.HF_JSON
+KEY_TOKENIZER_RWKV:       str = Keys.TOKENIZER.RWKV

gguf-py/gguf/gguf_writer.py

@@ -14,10 +14,10 @@
     GGUF_DEFAULT_ALIGNMENT,
     GGUF_MAGIC,
     GGUF_VERSION,
-    KEY,
     GGMLQuantizationType,
     GGUFEndian,
     GGUFValueType,
+    Keys,
     RopeScalingType,
     TokenType,
 )
@@ -278,132 +278,132 @@ def close(self) -> None:
         self.fout.close()
 
     def add_architecture(self) -> None:
-        self.add_string(KEY.GENERAL.ARCHITECTURE, self.arch)
+        self.add_string(Keys.GENERAL.ARCHITECTURE, self.arch)
 
     def add_author(self, author: str) -> None:
-        self.add_string(KEY.GENERAL.AUTHOR, author)
+        self.add_string(Keys.GENERAL.AUTHOR, author)
 
     def add_tensor_data_layout(self, layout: str) -> None:
-        self.add_string(KEY.LLM.TENSOR_DATA_LAYOUT.value.format(arch=self.arch), layout)
+        self.add_string(Keys.LLM.TENSOR_DATA_LAYOUT.value.format(arch=self.arch), layout)
 
     def add_url(self, url: str) -> None:
-        self.add_string(KEY.GENERAL.URL, url)
+        self.add_string(Keys.GENERAL.URL, url)
 
     def add_description(self, description: str) -> None:
-        self.add_string(KEY.GENERAL.DESCRIPTION, description)
+        self.add_string(Keys.GENERAL.DESCRIPTION, description)
 
     def add_source_url(self, url: str) -> None:
-        self.add_string(KEY.GENERAL.SOURCE_URL, url)
+        self.add_string(Keys.GENERAL.SOURCE_URL, url)
 
     def add_source_hf_repo(self, repo: str) -> None:
-        self.add_string(KEY.GENERAL.SOURCE_HF_REPO, repo)
+        self.add_string(Keys.GENERAL.SOURCE_HF_REPO, repo)
 
     def add_file_type(self, ftype: int) -> None:
-        self.add_uint32(KEY.GENERAL.FILE_TYPE, ftype)
+        self.add_uint32(Keys.GENERAL.FILE_TYPE, ftype)
 
     def add_name(self, name: str) -> None:
-        self.add_string(KEY.GENERAL.NAME, name)
+        self.add_string(Keys.GENERAL.NAME, name)
 
     def add_quantization_version(self, quantization_version: GGMLQuantizationType) -> None:
         self.add_uint32(
-            KEY.GENERAL.QUANTIZATION_VERSION, quantization_version)
+            Keys.GENERAL.QUANTIZATION_VERSION, quantization_version)
 
     def add_custom_alignment(self, alignment: int) -> None:
         self.data_alignment = alignment
-        self.add_uint32(KEY.GENERAL.ALIGNMENT, alignment)
+        self.add_uint32(Keys.GENERAL.ALIGNMENT, alignment)
 
     def add_context_length(self, length: int) -> None:
         self.add_uint32(
-            KEY.LLM.CONTEXT_LENGTH.value.format(arch=self.arch), length)
+            Keys.LLM.CONTEXT_LENGTH.value.format(arch=self.arch), length)
 
     def add_embedding_length(self, length: int) -> None:
         self.add_uint32(
-            KEY.LLM.EMBEDDING_LENGTH.value.format(arch=self.arch), length)
+            Keys.LLM.EMBEDDING_LENGTH.value.format(arch=self.arch), length)
 
     def add_block_count(self, length: int) -> None:
         self.add_uint32(
-            KEY.LLM.BLOCK_COUNT.value.format(arch=self.arch), length)
+            Keys.LLM.BLOCK_COUNT.value.format(arch=self.arch), length)
 
     def add_feed_forward_length(self, length: int) -> None:
         self.add_uint32(
-            KEY.LLM.FEED_FORWARD_LENGTH.value.format(arch=self.arch), length)
+            Keys.LLM.FEED_FORWARD_LENGTH.value.format(arch=self.arch), length)
 
     def add_parallel_residual(self, use: bool) -> None:
         self.add_bool(
-            KEY.LLM.USE_PARALLEL_RESIDUAL.value.format(arch=self.arch), use)
+            Keys.LLM.USE_PARALLEL_RESIDUAL.value.format(arch=self.arch), use)
 
     def add_head_count(self, count: int) -> None:
         self.add_uint32(
-            KEY.ATTENTION.HEAD_COUNT.value.format(arch=self.arch), count)
+            Keys.ATTENTION.HEAD_COUNT.value.format(arch=self.arch), count)
 
     def add_head_count_kv(self, count: int) -> None:
         self.add_uint32(
-            KEY.ATTENTION.HEAD_COUNT_KV.value.format(arch=self.arch), count)
+            Keys.ATTENTION.HEAD_COUNT_KV.value.format(arch=self.arch), count)
 
     def add_max_alibi_bias(self, bias: float) -> None:
         self.add_float32(
-            KEY.ATTENTION.MAX_ALIBI_BIAS.value.format(arch=self.arch), bias)
+            Keys.ATTENTION.MAX_ALIBI_BIAS.value.format(arch=self.arch), bias)
 
     def add_clamp_kqv(self, value: float) -> None:
         self.add_float32(
-            KEY.ATTENTION.CLAMP_KQV.value.format(arch=self.arch), value)
+            Keys.ATTENTION.CLAMP_KQV.value.format(arch=self.arch), value)
 
     def add_layer_norm_eps(self, value: float) -> None:
         self.add_float32(
-            KEY.ATTENTION.LAYERNORM_EPS.value.format(arch=self.arch), value)
+            Keys.ATTENTION.LAYERNORM_EPS.value.format(arch=self.arch), value)
 
     def add_layer_norm_rms_eps(self, value: float) -> None:
         self.add_float32(
-            KEY.ATTENTION.LAYERNORM_RMS_EPS.value.format(arch=self.arch), value)
+            Keys.ATTENTION.LAYERNORM_RMS_EPS.value.format(arch=self.arch), value)
 
     def add_rope_dimension_count(self, count: int) -> None:
         self.add_uint32(
-            KEY.ROPE.DIMENSION_COUNT.value.format(arch=self.arch), count)
+            Keys.ROPE.DIMENSION_COUNT.value.format(arch=self.arch), count)
 
     def add_rope_freq_base(self, value: float) -> None:
-        self.add_float32(KEY.ROPE.FREQ_BASE.value.format(arch=self.arch), value)
+        self.add_float32(Keys.ROPE.FREQ_BASE.value.format(arch=self.arch), value)
 
     def add_rope_scaling_type(self, value: RopeScalingType) -> None:
-        self.add_string(KEY.ROPE.SCALING_TYPE.value.format(arch=self.arch), value.value)
+        self.add_string(Keys.ROPE.SCALING_TYPE.value.format(arch=self.arch), value)
 
     def add_rope_scaling_factor(self, value: float) -> None:
-        self.add_float32(KEY.ROPE.SCALING_FACTOR.value.format(arch=self.arch), value)
+        self.add_float32(Keys.ROPE.SCALING_FACTOR.value.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.value.format(arch=self.arch), value)
+        self.add_uint32(Keys.ROPE.SCALING_ORIG_CTX_LEN.value.format(arch=self.arch), value)
 
     def add_rope_scaling_finetuned(self, value: bool) -> None:
-        self.add_bool(KEY.ROPE.SCALING_FINETUNED.value.format(arch=self.arch), value)
+        self.add_bool(Keys.ROPE.SCALING_FINETUNED.value.format(arch=self.arch), value)
 
     def add_tokenizer_model(self, model: str) -> None:
-        self.add_string(KEY.TOKENIZER.MODEL, model)
+        self.add_string(Keys.TOKENIZER.MODEL, model)
 
     def add_token_list(self, tokens: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
-        self.add_array(KEY.TOKENIZER.LIST, tokens)
+        self.add_array(Keys.TOKENIZER.LIST, tokens)
 
     def add_token_merges(self, merges: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
-        self.add_array(KEY.TOKENIZER.MERGES, merges)
+        self.add_array(Keys.TOKENIZER.MERGES, merges)
 
     def add_token_types(self, types: Sequence[TokenType] | Sequence[int]) -> None:
-        self.add_array(KEY.TOKENIZER.TOKEN_TYPE, types)
+        self.add_array(Keys.TOKENIZER.TOKEN_TYPE, types)
 
     def add_token_scores(self, scores: Sequence[float]) -> None:
-        self.add_array(KEY.TOKENIZER.SCORES, scores)
+        self.add_array(Keys.TOKENIZER.SCORES, scores)
 
     def add_bos_token_id(self, id: int) -> None:
-        self.add_uint32(KEY.TOKENIZER.BOS_ID, id)
+        self.add_uint32(Keys.TOKENIZER.BOS_ID, id)
 
     def add_eos_token_id(self, id: int) -> None:
-        self.add_uint32(KEY.TOKENIZER.EOS_ID, id)
+        self.add_uint32(Keys.TOKENIZER.EOS_ID, id)
 
     def add_unk_token_id(self, id: int) -> None:
-        self.add_uint32(KEY.TOKENIZER.UNK_ID, id)
+        self.add_uint32(Keys.TOKENIZER.UNK_ID, id)
 
     def add_sep_token_id(self, id: int) -> None:
-        self.add_uint32(KEY.TOKENIZER.SEP_ID, id)
+        self.add_uint32(Keys.TOKENIZER.SEP_ID, id)
 
     def add_pad_token_id(self, id: int) -> None:
-        self.add_uint32(KEY.TOKENIZER.PAD_ID, id)
+        self.add_uint32(Keys.TOKENIZER.PAD_ID, id)
 
     def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes:
         pack_prefix = ''