• Matt's avatar
    Support sharded safetensors in TF (#29350) · 11ef35e8
    Matt authored
    
    
    * Initial commit (still lots of unfinished bits)
    
    * (Still untested) add safetensors sharding to save_pretrained
    
    * Fix savetensors saving, update default shard size to match PT
    
    * Add proper loading of TF-format safetensors
    
    * Revert default size in case that changes things
    
    * Fix incorrect index name
    
    * Update loading priority
    
    * Update tests
    
    * Make the tests a little more stringent
    
    * Expand tests
    
    * Add sharded cross-test
    
    * Fix argument name
    
    * One more test fix
    
    * Adding mlx to the list of allowed formats
    
    * Remove irrelevant block for safetensors
    
    * Refactor warning logging into a separate function
    
    * Remove unused skip_logger_warnings arg
    
    * Update src/transformers/modeling_tf_utils.py
    Co-authored-by: default avataramyeroberts <22614925+amyeroberts@users.noreply.github.com>
    
    * Move function def
    
    ---------
    Co-authored-by: default avataramyeroberts <22614925+amyeroberts@users.noreply.github.com>
    11ef35e8
test_modeling_tf_utils.py 35.4 KB