From 279a9ba1cb516948bf71754f973bfc2d7868fee2 Mon Sep 17 00:00:00 2001
From: Kabir Brar <kabir@brar.org>
Date: Tue, 25 Jul 2023 10:05:33 -0700
Subject: [PATCH] int: Fix quantization schema (#3479)

---
 ludwig/schema/llms/quantization.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/ludwig/schema/llms/quantization.py b/ludwig/schema/llms/quantization.py
index 04eae3380b3..2812f59afa9 100644
--- 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": "",
+        }