@@ -43,17 +43,18 @@ class SentencePieceTokenTypes(IntEnum):
4343class Model (ABC ):
4444 _model_classes : dict [str , type [Model ]] = {}
4545
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 ):
4747 self .dir_model = dir_model
4848 self .ftype = ftype
4949 self .fname_out = fname_out
5050 self .is_big_endian = is_big_endian
5151 self .endianess = gguf .GGUFEndian .BIG if is_big_endian else gguf .GGUFEndian .LITTLE
52+ self .use_temp_file = use_temp_file
5253 self .is_safetensors = self ._is_model_safetensors ()
5354 self .num_parts = Model .count_model_parts (self .dir_model , ".safetensors" if self .is_safetensors else ".bin" )
5455 self .part_names = self ._get_part_names ()
5556 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 )
5758 self .block_count = self .find_hparam (["n_layers" , "num_hidden_layers" , "n_layer" ])
5859
5960 @property
@@ -2459,6 +2460,7 @@ def parse_args() -> argparse.Namespace:
24592460 "model" , type = Path ,
24602461 help = "directory containing model file" ,
24612462 )
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)" )
24622464
24632465 return parser .parse_args ()
24642466
@@ -2502,7 +2504,7 @@ def main() -> None:
25022504
25032505 with torch .inference_mode ():
25042506 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 )
25062508
25072509 print ("Set model parameters" )
25082510 model_instance .set_gguf_parameters ()
0 commit comments