move kv keys to constants.py

This commit is contained in:
Christian Zhou-Zheng 2024-06-06 08:49:46 -04:00
parent 1cbab22225
commit 2037eabb64
2 changed files with 9 additions and 7 deletions

View file

@ -72,6 +72,11 @@ class Keys:
SCALING_FINETUNED = "{arch}.rope.scaling.finetuned" SCALING_FINETUNED = "{arch}.rope.scaling.finetuned"
SCALING_YARN_LOG_MUL = "{arch}.rope.scaling.yarn_log_multiplier" 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: class SSM:
CONV_KERNEL = "{arch}.ssm.conv_kernel" CONV_KERNEL = "{arch}.ssm.conv_kernel"
INNER_SIZE = "{arch}.ssm.inner_size" INNER_SIZE = "{arch}.ssm.inner_size"

View file

@ -18,15 +18,12 @@ from .constants import (
GGUFValueType GGUFValueType
) )
from .gguf_writer import GGUFWriter, WriterState from .gguf_writer import GGUFWriter, WriterState
from .constants import Keys
SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf" SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf"
METADATA_ONLY_INDICATOR = -1 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)} 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) 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 # 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: 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(Keys.Split.LLM_KV_SPLIT_NO, i)
writer.add_uint16(LLM_KV_SPLIT_COUNT, len(self.shards)) writer.add_uint16(Keys.Split.LLM_KV_SPLIT_COUNT, len(self.shards))
writer.add_int32(LLM_KV_SPLIT_TENSORS_COUNT, self.total_tensors) 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 # add tensors, deque popleft() ensures references to eager tensors are not kept
while True: while True: