Skip to content
This repository was archived by the owner on Jun 3, 2025. It is now read-only.

Commit cc0900a

Browse files
authored
Improve KVCacheInjector config failure message (#1687)
1 parent ab16210 commit cc0900a

File tree

1 file changed

+6
-4
lines changed

1 file changed

+6
-4
lines changed

src/sparseml/exporters/kv_cache_injector.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -75,8 +75,8 @@ def __init__(
7575
Alternatively:
7676
```python
7777
onnx_model: onnx.ModelProto = ...
78-
exporter = KeyValueCacheInjector(model_path="path/to/model")
79-
exporter = KeyValueCacheInjector(num_attention_heads = 16,
78+
exporter = KeyValueCacheInjector(model_path="path/to/model",
79+
num_attention_heads = 16,
8080
hidden_size_dim = 64)
8181
exporter.export(onnx_model, "model.onnx")
8282
```
@@ -106,8 +106,10 @@ def __init__(
106106

107107
else:
108108
raise ValueError(
109-
"Either `model_path` or kwargs must be provided to "
110-
"KeyValueCacheInjector"
109+
f"Unable to find KeyValueCacheConfig for model_path='{model_path}'. "
110+
"Either kwargs must be provided to KeyValueCacheInjector to construct "
111+
"OnnxTransform, or a new config should be registered in "
112+
"`sparseml/src/sparseml/exporters/transforms/kv_cache/configs.py`"
111113
)
112114

113115
super().__init__(transforms)

0 commit comments

Comments
 (0)