Add input validation for structured-data-classification (#97)

Author: osanseviero

api-inference-community/api_inference_community/validation.py

@@ -132,6 +132,18 @@ def all_rows_must_have_same_length(cls, table: Dict[str, List[str]]):
         raise ValueError("All rows in the table must be the same length")
 
 
+class StructuredDataClassificationInputsCheck(BaseModel):
+    data: Dict[str, List[str]]
+
+    @validator("data")
+    def all_rows_must_have_same_length(cls, data: Dict[str, List[str]]):
+        rows = list(data.values())
+        n = len(rows[0])
+        if all(len(x) == n for x in rows):
+            return data
+        raise ValueError("All rows in the data must be the same length")
+
+
 class StringOrStringBatchInputCheck(BaseModel):
     __root__: Union[List[str], str]
 
@@ -164,6 +176,7 @@ class StringInput(BaseModel):
     "feature-extraction": StringOrStringBatchInputCheck,
     "sentence-similarity": SentenceSimilarityInputsCheck,
     "table-question-answering": TableQuestionAnsweringInputsCheck,
+    "structured-data-classification": StructuredDataClassificationInputsCheck,
     "fill-mask": StringInput,
     "summarization": StringInput,
     "text2text-generation": StringInput,

api-inference-community/tests/test_nlp.py

@@ -152,6 +152,39 @@ def test_no_query(self):
             normalize_payload_nlp(bpayload, "table-question-answering")
 
 
+class StructuredDataClassificationValidationTestCase(TestCase):
+    def test_valid_input(self):
+        data = {
+            "Repository": ["Transformers", "Datasets", "Tokenizers"],
+            "Stars": ["36542", "4512", "3934"],
+        }
+
+        inputs = {"data": data}
+        bpayload = json.dumps({"inputs": inputs}).encode("utf-8")
+        normalized_inputs, processed_params = normalize_payload_nlp(
+            bpayload, "structured-data-classification"
+        )
+        self.assertEqual(processed_params, {})
+        self.assertEqual(inputs, normalized_inputs)
+
+    def test_invalid_data_lengths(self):
+        data = {
+            "Repository": ["Transformers", "Datasets", "Tokenizers"],
+            "Stars": ["36542", "4512"],
+        }
+
+        inputs = {"data": data}
+        bpayload = json.dumps({"inputs": inputs}).encode("utf-8")
+        with self.assertRaises(ValidationError):
+            normalize_payload_nlp(bpayload, "structured-data-classification")
+
+    def test_invalid_data_type(self):
+        inputs = {"data": "Invalid data"}
+        bpayload = json.dumps({"inputs": inputs}).encode("utf-8")
+        with self.assertRaises(ValidationError):
+            normalize_payload_nlp(bpayload, "structured-data-classification")
+
+
 class SummarizationValidationTestCase(TestCase):
     def test_no_params(self):
         bpayload = json.dumps({"inputs": "whatever"}).encode("utf-8")