int: Fix quantization schema (#3479)

ludwig-ai · Jul 25, 2023 · 279a9ba · 279a9ba
1 parent 48f7d1f
commit 279a9ba
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/ludwig/schema/llms/quantization.py b/ludwig/schema/llms/quantization.py
@@ -86,4 +86,15 @@ def __init__(self):
         super().__init__(QuantizationConfig, default_missing=True)
 
     def _jsonschema_type_mapping(self):
-        return schema_utils.unload_jsonschema_from_marshmallow_class(QuantizationConfig)
+        return {
+            "oneOf": [
+                {"type": "null", "title": "disabled", "description": "Disable quantization."},
+                {
+                    **schema_utils.unload_jsonschema_from_marshmallow_class(QuantizationConfig),
+                    "title": "enabled",
+                    "description": "Set quantization options.",
+                },
+            ],
+            "title": "quantization",
+            "description": "",
+        }