@@ -43,17 +43,18 @@ class SentencePieceTokenTypes(IntEnum):
43
43
class Model (ABC ):
44
44
_model_classes : dict [str , type [Model ]] = {}
45
45
46
- def __init__ (self , dir_model : Path , ftype : int , fname_out : Path , is_big_endian : bool ):
46
+ def __init__ (self , dir_model : Path , ftype : int , fname_out : Path , is_big_endian : bool , use_temp_file : bool ):
47
47
self .dir_model = dir_model
48
48
self .ftype = ftype
49
49
self .fname_out = fname_out
50
50
self .is_big_endian = is_big_endian
51
51
self .endianess = gguf .GGUFEndian .BIG if is_big_endian else gguf .GGUFEndian .LITTLE
52
+ self .use_temp_file = use_temp_file
52
53
self .is_safetensors = self ._is_model_safetensors ()
53
54
self .num_parts = Model .count_model_parts (self .dir_model , ".safetensors" if self .is_safetensors else ".bin" )
54
55
self .part_names = self ._get_part_names ()
55
56
self .hparams = Model .load_hparams (self .dir_model )
56
- self .gguf_writer = gguf .GGUFWriter (fname_out , gguf .MODEL_ARCH_NAMES [self .model_arch ], endianess = self .endianess , use_temp_file = False )
57
+ self .gguf_writer = gguf .GGUFWriter (fname_out , gguf .MODEL_ARCH_NAMES [self .model_arch ], endianess = self .endianess , use_temp_file = self . use_temp_file )
57
58
self .block_count = self .find_hparam (["n_layers" , "num_hidden_layers" , "n_layer" ])
58
59
59
60
@property
@@ -2459,6 +2460,7 @@ def parse_args() -> argparse.Namespace:
2459
2460
"model" , type = Path ,
2460
2461
help = "directory containing model file" ,
2461
2462
)
2463
+ parser .add_argument ("--use-temp-file" , action = "store_true" , help = "use the tempfile library while processing (helpful when running out of memory, process killed)" )
2462
2464
2463
2465
return parser .parse_args ()
2464
2466
@@ -2502,7 +2504,7 @@ def main() -> None:
2502
2504
2503
2505
with torch .inference_mode ():
2504
2506
model_class = Model .from_model_architecture (hparams ["architectures" ][0 ])
2505
- model_instance = model_class (dir_model , ftype_map [args .outtype ], fname_out , args .bigendian )
2507
+ model_instance = model_class (dir_model , ftype_map [args .outtype ], fname_out , args .bigendian , args . use_temp_file )
2506
2508
2507
2509
print ("Set model parameters" )
2508
2510
model_instance .set_gguf_parameters ()
0 commit comments