diff --git a/src/outlines/types/dsl.py b/src/outlines/types/dsl.py index 5949a67a4..19504c392 100644 --- a/src/outlines/types/dsl.py +++ b/src/outlines/types/dsl.py @@ -574,6 +574,10 @@ class QuantifyExact(Term): term: Term count: int + def __post_init__(self): + if self.count < 0: + raise ValueError("QuantifyExact: `count` must be a non-negative integer.") + def _display_node(self) -> str: return f"Quantify({{{self.count}}})" @@ -589,6 +593,12 @@ class QuantifyMinimum(Term): term: Term min_count: int + def __post_init__(self): + if self.min_count < 0: + raise ValueError( + "QuantifyMinimum: `min_count` must be a non-negative integer." + ) + def _display_node(self) -> str: return f"Quantify({{{self.min_count},}})" @@ -606,6 +616,12 @@ class QuantifyMaximum(Term): term: Term max_count: int + def __post_init__(self): + if self.max_count < 0: + raise ValueError( + "QuantifyMaximum: `max_count` must be a non-negative integer." + ) + def _display_node(self) -> str: return f"Quantify({{,{self.max_count}}})" @@ -625,6 +641,10 @@ class QuantifyBetween(Term): max_count: int def __post_init__(self): + if self.min_count < 0: + raise ValueError( + "QuantifyBetween: `min_count` must be a non-negative integer." + ) if self.min_count > self.max_count: raise ValueError( "QuantifyBetween: `max_count` must be greater than `min_count`." diff --git a/tests/types/test_dsl.py b/tests/types/test_dsl.py index f827d9386..94938325c 100644 --- a/tests/types/test_dsl.py +++ b/tests/types/test_dsl.py @@ -138,6 +138,15 @@ def test_dsl_init(): assert repr(maximum) == "QuantifyMaximum(term=String(value='test'), max_count=3)" assert maximum.display_ascii_tree() == "└── Quantify({,3})\n └── String('test')\n" + with pytest.raises(ValueError, match="`count` must be a non-negative integer"): + QuantifyExact(string, -1) + + with pytest.raises(ValueError, match="`min_count` must be a non-negative integer"): + QuantifyMinimum(string, -1) + + with pytest.raises(ValueError, match="`max_count` must be a non-negative integer"): + QuantifyMaximum(string, -1) + between = QuantifyBetween(string, 1, 3) assert between.term == string assert between.min_count == 1 @@ -153,6 +162,9 @@ def test_dsl_init(): ): QuantifyBetween(string, 3, 1) + with pytest.raises(ValueError, match="`min_count` must be a non-negative integer"): + QuantifyBetween(string, -2, -1) + def test_dsl_term_methods(): a = String("a")