From f4837ed253862f6742ffed2bdebbc9d398f9df35 Mon Sep 17 00:00:00 2001 From: Sara Adkins Date: Wed, 10 Apr 2024 19:42:16 +0000 Subject: [PATCH] bump pydantic --- setup.py | 2 +- src/sparsetensors/README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index d3a6f905..12ef67a3 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,7 @@ def _setup_packages() -> List: ) def _setup_install_requires() -> List: - return ["torch>=1.7.0", "transformers<=4.40", "pydantic>=1.8.2,<2.0.0"] + return ["torch>=1.7.0", "transformers<=4.40", "pydantic<2.7"] def _setup_extras() -> Dict: return {"dev": ["black==22.12.0", "isort==5.8.0", "wheel>=0.36.2", "flake8>=3.8.3", "pytest>=6.0.0"]} diff --git a/src/sparsetensors/README.md b/src/sparsetensors/README.md index 51d49ade..833d1ec2 100644 --- a/src/sparsetensors/README.md +++ b/src/sparsetensors/README.md @@ -65,7 +65,7 @@ is saved to the model's `config.json` file. **Note:** the model must have been initialized with SparseAutoModelForCausalLM.from_pretrained() ```python -from sparseml.transformers.compression import BitmaskConfig +from sparsetensors import BitmaskConfig output_dir = "/PATH/TO/SAVE/COMPRESSED_MODEL" sparsity_config = BitmaskConfig()