Small refactor
This commit is contained in:
parent
c94df09732
commit
235acc18cd
2 changed files with 23 additions and 27 deletions
|
@ -28,41 +28,25 @@ else:
|
||||||
# output in the same directory as the model by default
|
# output in the same directory as the model by default
|
||||||
fname_out = dir_model / f'ggml-model-{ftype_str[ftype]}.gguf'
|
fname_out = dir_model / f'ggml-model-{ftype_str[ftype]}.gguf'
|
||||||
|
|
||||||
|
print(f"Loading model: {dir_model.name}")
|
||||||
print("gguf: loading model " + dir_model.name)
|
|
||||||
|
|
||||||
hparams = model.Model.load_hparams(dir_model)
|
hparams = model.Model.load_hparams(dir_model)
|
||||||
|
|
||||||
model_class = model.Model.from_model_architecture(hparams["architectures"][0])
|
model_class = model.Model.from_model_architecture(hparams["architectures"][0])
|
||||||
model_instance = model_class(dir_model, ftype, fname_out)
|
model_instance = model_class(dir_model, ftype, fname_out)
|
||||||
|
|
||||||
print("gguf: get model metadata")
|
print("Set model parameters")
|
||||||
|
|
||||||
model_instance.set_gguf_parameters()
|
model_instance.set_gguf_parameters()
|
||||||
|
|
||||||
# TOKENIZATION
|
print("Set model tokenizer")
|
||||||
print("gguf: get tokenizer metadata")
|
|
||||||
|
|
||||||
print("gguf: get gpt2 tokenizer vocab")
|
|
||||||
|
|
||||||
model_instance.set_vocab()
|
model_instance.set_vocab()
|
||||||
|
|
||||||
# write model
|
|
||||||
if not args.vocab_only:
|
if not args.vocab_only:
|
||||||
model_instance.write_tensors()
|
print(f"Exporting model to '{fname_out}'")
|
||||||
print("gguf: write header")
|
model_instance.write()
|
||||||
model_instance.gguf_writer.write_header_to_file()
|
|
||||||
print("gguf: write metadata")
|
|
||||||
model_instance.gguf_writer.write_kv_data_to_file()
|
|
||||||
print("gguf: write tensors")
|
|
||||||
model_instance.gguf_writer.write_tensors_to_file()
|
|
||||||
else:
|
else:
|
||||||
print("gguf: write header")
|
print(f"Exporting model vocab to '{fname_out}'")
|
||||||
model_instance.gguf_writer.write_header_to_file()
|
model_instance.write_vocab()
|
||||||
print("gguf: write metadata")
|
|
||||||
model_instance.gguf_writer.write_kv_data_to_file()
|
|
||||||
|
|
||||||
model_instance.gguf_writer.close()
|
print(f"Model successfully exported to '{fname_out}'")
|
||||||
|
|
||||||
print(f"gguf: model successfully exported to '{fname_out}'")
|
|
||||||
print("")
|
print("")
|
||||||
|
|
18
model.py
18
model.py
|
@ -9,7 +9,7 @@ import numpy as np
|
||||||
|
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TypeAlias, Any
|
from typing import TypeAlias, Any, Generator
|
||||||
|
|
||||||
NDArray: TypeAlias = 'np.ndarray[Any, Any]'
|
NDArray: TypeAlias = 'np.ndarray[Any, Any]'
|
||||||
|
|
||||||
|
@ -48,7 +48,7 @@ class Model:
|
||||||
return ("pytorch_model.bin",)
|
return ("pytorch_model.bin",)
|
||||||
return (f"pytorch_model-{n:05}-of-{self.num_parts:05}.bin" for n in range(1, self.num_parts + 1))
|
return (f"pytorch_model-{n:05}-of-{self.num_parts:05}.bin" for n in range(1, self.num_parts + 1))
|
||||||
|
|
||||||
def _get_model_architecture(self):
|
def _get_model_architecture(self) -> gguf.MODEL_ARCH:
|
||||||
arch = self.hparams["architectures"][0]
|
arch = self.hparams["architectures"][0]
|
||||||
if arch == "GPTNeoXForCausalLM":
|
if arch == "GPTNeoXForCausalLM":
|
||||||
return gguf.MODEL_ARCH.GPTNEOX
|
return gguf.MODEL_ARCH.GPTNEOX
|
||||||
|
@ -160,7 +160,7 @@ class Model:
|
||||||
def set_vocab(self):
|
def set_vocab(self):
|
||||||
self._set_vocab_gpt2()
|
self._set_vocab_gpt2()
|
||||||
|
|
||||||
def get_tensors(self):
|
def get_tensors(self) -> Generator[str, Any]:
|
||||||
for part_name in self.part_names:
|
for part_name in self.part_names:
|
||||||
print("gguf: loading model part '" + part_name + "'")
|
print("gguf: loading model part '" + part_name + "'")
|
||||||
if self.is_safetensors:
|
if self.is_safetensors:
|
||||||
|
@ -230,6 +230,18 @@ class Model:
|
||||||
|
|
||||||
self.gguf_writer.add_tensor(new_name, data)
|
self.gguf_writer.add_tensor(new_name, data)
|
||||||
|
|
||||||
|
def write(self):
|
||||||
|
self.write_tensors()
|
||||||
|
self.gguf_writer.write_header_to_file()
|
||||||
|
self.gguf_writer.write_kv_data_to_file()
|
||||||
|
self.gguf_writer.write_tensors_to_file()
|
||||||
|
self.gguf_writer.close()
|
||||||
|
|
||||||
|
def write_vocab(self):
|
||||||
|
self.gguf_writer.write_header_to_file()
|
||||||
|
self.gguf_writer.write_kv_data_to_file()
|
||||||
|
self.gguf_writer.close()
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def count_model_parts(dir_model: Path, prefix: str) -> int:
|
def count_model_parts(dir_model: Path, prefix: str) -> int:
|
||||||
num_parts = 0
|
num_parts = 0
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue