File tree Expand file tree Collapse file tree 2 files changed +25
-4
lines changed
transformers/sparsification Expand file tree Collapse file tree 2 files changed +25
-4
lines changed Original file line number Diff line number Diff line change 2727def save_checkpoint (
2828 save_path : str ,
2929 model : PreTrainedModel ,
30- processor : Processor ,
30+ processor : Optional [ Processor ] = None ,
3131 save_safetensors : bool = True ,
3232 save_compressed : bool = True ,
3333 skip_sparsity_compression_stats : bool = False ,
3434):
35+ """
36+ Save a model, processor, and recipe
37+
38+ :param save_path: Path used to save model and processor
39+ :param model: model to save
40+ :param processor: processor to save
41+ :param save_safetensors: save model checkpoint using safetensors file type
42+ :param save_compressed: save model checkpoint using compressed-tensors format
43+ """
3544 # saving the model also saves the recipe
3645 model .save_pretrained (
3746 save_path ,
Original file line number Diff line number Diff line change 3636def modify_save_pretrained (model : PreTrainedModel ):
3737 """
3838 Overrides a PreTrainedModel's save_pretrained() method with a wrapped version that
39- supports compression
39+ supports compression. The new save_pretrained function performs the following saving
40+ operations:
41+
42+ 1. Saves the model state, potentially in a compressed format
43+ 2. Saves the recipe, appending any current recipes to existing recipe files
44+ 3. Copies any necessary python files from the model cache
4045 """
4146
4247 def save_pretrained_compressed (save_pretrained_method ):
@@ -279,9 +284,16 @@ def get_model_compressor(
279284 )
280285
281286
282- def update_and_save_recipe (model_path : str , save_directory : str ):
287+ def update_and_save_recipe (model_stub : str , save_directory : str ):
288+ """
289+ Save a recipe ontop of any existing recipe files located at model_stub
290+
291+ :param model_stub: path to existing model or model stub which may contain an
292+ existing recipe
293+ :param save_directory: path to save combined existing recipe and current recipe
294+ """
283295 recipes_to_save = []
284- existing_recipe = infer_recipe_from_model_path (model_path )
296+ existing_recipe = infer_recipe_from_model_path (model_stub )
285297 if existing_recipe is not None :
286298 recipes_to_save .append (existing_recipe )
287299
You can’t perform that action at this time.
0 commit comments