From 2037eabb644873400f6cdf0514f66af2a10076d2 Mon Sep 17 00:00:00 2001 From: Christian Zhou-Zheng Date: Thu, 6 Jun 2024 08:49:46 -0400 Subject: [PATCH] move kv keys to constants.py --- gguf-py/gguf/constants.py | 5 +++++ gguf-py/gguf/gguf_manager.py | 11 ++++------- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a3c024c89..a5bab4de6 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -72,6 +72,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" diff --git a/gguf-py/gguf/gguf_manager.py b/gguf-py/gguf/gguf_manager.py index 523a5f500..f4411e752 100644 --- a/gguf-py/gguf/gguf_manager.py +++ b/gguf-py/gguf/gguf_manager.py @@ -18,15 +18,12 @@ from .constants import ( GGUFValueType ) from .gguf_writer import GGUFWriter, WriterState +from .constants import Keys SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf" METADATA_ONLY_INDICATOR = -1 -LLM_KV_SPLIT_NO = "split.no" -LLM_KV_SPLIT_COUNT = "split.count" -LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count" - KVTempData: TypeAlias = dict[str, tuple[Any, GGUFValueType]] # {key: (value, type)} TensorTempData: TypeAlias = tuple[str, np.ndarray[Any, Any], GGMLQuantizationType] # (tensor name, tensor data, tensor dtype) @@ -132,9 +129,9 @@ class GGUFManager(GGUFWriter): # add split metadata unless it's one file - small first shard splits even with SplitStyle.NONE if self.split_arguments.split_style != SplitStyle.NONE or self.split_arguments.small_first_shard: - writer.add_uint16(LLM_KV_SPLIT_NO, i) - writer.add_uint16(LLM_KV_SPLIT_COUNT, len(self.shards)) - writer.add_int32(LLM_KV_SPLIT_TENSORS_COUNT, self.total_tensors) + writer.add_uint16(Keys.Split.LLM_KV_SPLIT_NO, i) + writer.add_uint16(Keys.Split.LLM_KV_SPLIT_COUNT, len(self.shards)) + writer.add_int32(Keys.Split.LLM_KV_SPLIT_TENSORS_COUNT, self.total_tensors) # add tensors, deque popleft() ensures references to eager tensors are not kept while True: