update tests

dsikka · dsikka · commit ca5279ff3715 · 2024-06-24T14:43:55.000Z
diff --git a/tests/test_compressors/test_pack_quant.py b/tests/test_compressors/test_pack_quant.py
@@ -67,7 +67,7 @@ def test_quant_format(shape):
     compressor = PackedQuantizationCompressor(config=quant_config)
     quantized_modules_to_args = {"dummy": quant_config.config_groups["group_1"].weights}
     compressed_state_dict = compressor.compress(
-        dense_state_dict, model_quant_args=quantized_modules_to_args
+        dense_state_dict, names_to_scheme=quantized_modules_to_args
     )
 
     # compressed state_dict adds one entry for shape
@@ -129,7 +129,7 @@ def test_reload_match(tmp_path, num_bits):
         "dummy2": quant_config.config_groups["group_1"].weights,
     }
     compressed_state_dict = compressor.compress(
-        dense_state_dict, model_quant_args=quantized_modules_to_args
+        dense_state_dict, names_to_scheme=quantized_modules_to_args
     )
     save_file(compressed_state_dict, tmp_path / "model.safetensors")
     reconstructed_dense_gen = compressor.decompress(

Original file line number	Diff line number	Diff line change
`@@ -67,7 +67,7 @@ def test_quant_format(shape):`
`67`	`67`	`compressor = PackedQuantizationCompressor(config=quant_config)`
`68`	`68`	`quantized_modules_to_args = {"dummy": quant_config.config_groups["group_1"].weights}`
`69`	`69`	`compressed_state_dict = compressor.compress(`
`70`		`- dense_state_dict, model_quant_args=quantized_modules_to_args`
	`70`	`+ dense_state_dict, names_to_scheme=quantized_modules_to_args`
`71`	`71`	`)`
`72`	`72`
`73`	`73`	`# compressed state_dict adds one entry for shape`
`@@ -129,7 +129,7 @@ def test_reload_match(tmp_path, num_bits):`
`129`	`129`	`"dummy2": quant_config.config_groups["group_1"].weights,`
`130`	`130`	`}`
`131`	`131`	`compressed_state_dict = compressor.compress(`
`132`		`- dense_state_dict, model_quant_args=quantized_modules_to_args`
	`132`	`+ dense_state_dict, names_to_scheme=quantized_modules_to_args`
`133`	`133`	`)`
`134`	`134`	`save_file(compressed_state_dict, tmp_path / "model.safetensors")`
`135`	`135`	`reconstructed_dense_gen = compressor.decompress(`