Skip to content

Commit

Permalink
attempt to appease the linter
Browse files Browse the repository at this point in the history
  • Loading branch information
christianazinn committed Jun 8, 2024
1 parent c6ae1d6 commit 2e70fa1
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 17 deletions.
2 changes: 1 addition & 1 deletion convert-hf-to-gguf.py
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,7 @@ def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path,
# allow templating the file name with the output ftype, useful with the "auto" ftype
self.fname_out = fname_out.parent / fname_out.name.format(ftype_lw, outtype=ftype_lw, ftype=ftype_lw, OUTTYPE=ftype_up, FTYPE=ftype_up)
self.gguf_writer = gguf.GGUFWriterSplit(fname_out, gguf.MODEL_ARCH_NAMES[self.model_arch], split_arguments,
endianess=self.endianess, use_temp_file=self.use_temp_file)
endianess=self.endianess, use_temp_file=self.use_temp_file)

@classmethod
def __init_subclass__(cls):
Expand Down
33 changes: 17 additions & 16 deletions gguf-py/gguf/gguf_writer_split.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import os
import logging
from enum import IntEnum
from typing import TYPE_CHECKING, Any, Sequence
from argparse import Namespace
Expand All @@ -21,6 +22,8 @@
from .gguf_writer import GGUFWriter, WriterState
from .constants import Keys

logger = logging.getLogger(__name__)


SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf"
METADATA_ONLY_INDICATOR = -1
Expand Down Expand Up @@ -63,7 +66,7 @@ class GGUFWriterSplit(GGUFWriter):

def __init__(self, path: os.PathLike[str] | str, arch: str, split_arguments: SplitArguments,
use_temp_file: bool = True, endianess: GGUFEndian = GGUFEndian.LITTLE
) -> None:
) -> None:
# we intentionally don't call superclass constructor
self.arch = arch
self.path = Path(path)
Expand All @@ -86,11 +89,11 @@ def init_shards(self) -> None:

# check if we need to split
if self.split_arguments.split_max_tensors and self.total_tensors < self.split_arguments.split_max_tensors:
print("Model has fewer tensors than the split threshold, not splitting")
logger.warning("Model has fewer tensors than the split threshold, not splitting")
self.split_style = SplitStyle.NONE

if self.split_arguments.split_max_size and total_size < self.split_arguments.split_max_size:
print("Model has smaller size than the split threshold, not splitting")
logger.warning("Model has smaller size than the split threshold, not splitting")
self.split_style = SplitStyle.NONE

# no shards are created when writing vocab so make one
Expand All @@ -105,13 +108,12 @@ def init_shards(self) -> None:
self.shards[i].path = self.path.with_name(SHARD_NAME_FORMAT.format(self.path.stem, i + 1, len(self.shards)))

# print shard info
print("\nWriting the following files:")
logger.info("Writing the following files:")
for shard in self.shards:
print(f" {shard.path}: n_tensors = {shard.tensor_count}, total_size = {GGUFWriterSplit.format_n_bytes_to_str(shard.size)}")
print()
logger.info(f" {shard.path}: n_tensors = {shard.tensor_count}, total_size = {GGUFWriterSplit.format_n_bytes_to_str(shard.size)}")

if self.split_arguments.dry_run:
print("\nDry run, not writing files")
logger.info("Dry run, not writing files")
exit()

# we don't want to initialize GGUFWriters until now because they create files
Expand All @@ -137,7 +139,7 @@ def init_shards(self) -> None:
try:
(name, tensor, dtype) = shard.tensors.popleft()
writer.add_tensor(name, tensor, raw_dtype=dtype)
except:
except IndexError:
break

self.shard_writers.append(writer)
Expand All @@ -154,7 +156,7 @@ def write_header_to_file(self) -> None:
def write_kv_data_to_file(self) -> None:
if self.state is not WriterState.HEADER:
raise ValueError(f'Expected GGUFWriterSplit state to be HEADER, got {self.state}')

for writer in self.shard_writers:
writer.write_kv_data_to_file()

Expand All @@ -169,9 +171,9 @@ def write_tensors_to_file(self, *, progress: bool = False) -> None:
writer = self.shard_writers[i]
is_metadata = writer.ti_data_count == 0
if is_metadata:
print(f"Writing to shard {i + 1}/{len(self.shards)} with metadata only")
logger.info(f"Writing to shard {i + 1}/{len(self.shards)} with metadata only")
else:
print(f"Writing to shard {i + 1}/{len(self.shards)} with {writer.ti_data_count}/{running_total} remaining tensors (of {self.total_tensors} total)")
logger.info(f"Writing to shard {i + 1}/{len(self.shards)} with {writer.ti_data_count}/{running_total} remaining tensors (of {self.total_tensors} total)")
running_total -= writer.ti_data_count
writer.write_tensors_to_file(progress=(progress and not is_metadata))
del writer
Expand All @@ -181,7 +183,7 @@ def write_tensors_to_file(self, *, progress: bool = False) -> None:
# override add_key, add_val to handle kv data separately
def add_key(self, key: str) -> None:
self.recent_key = key

def add_val(self, val: Any, vtype: GGUFValueType | None = None, add_vtype: bool = True) -> None:
if self.recent_key is None:
raise ValueError("No key set for value")
Expand Down Expand Up @@ -226,9 +228,7 @@ def get_tensor_size(tensor) -> int:
return tensor.data_type.elements_to_bytes(np.prod(tensor.shape))
except AttributeError: # numpy ndarray[Any, Any]
return tensor.nbytes
except: # this should never happen
raise ValueError(f"Invalid tensor type: {type(tensor)}")


@staticmethod
def split_str_to_n_bytes(split_str: str) -> int:
if split_str.endswith("K"):
Expand Down Expand Up @@ -256,4 +256,5 @@ def format_n_bytes_to_str(num: int) -> str:
if abs(fnum) < 1000.0:
return f"{fnum:3.1f}{unit}"
fnum /= 1000.0
return f"{fnum:.1f}T - over 1TB, --split recommended"
return f"{fnum:.1f}T - over 1TB, --split recommended"

0 comments on commit 2e70fa1

Please sign in to comment.