better errors and handling of data types

Author: isabelizimm

vetiver/handlers/spacy.py

@@ -33,16 +33,24 @@ def construct_prototype(self):
         prototype :
             Input data prototype for spacy model
         """
-        if self.prototype_data is None:
+        if self.prototype_data is not None and not isinstance(
+            self.prototype_data, (pd.Series, pd.DataFrame, dict)
+        ):  # wrong type
+            raise TypeError(
+                "Spacy prototype must be a dict, pandas Series, or pandas DataFrame"
+            )
+        elif (
+            isinstance(self.prototype_data, pd.DataFrame)
+            and len(self.prototype_data.columns) != 1
+        ):  # is dataframe, more than one column
+            raise ValueError("Spacy prototype data must be a 1-column pandas DataFrame")
+        elif (
+            isinstance(self.prototype_data, dict) and len(self.prototype_data) != 1
+        ):  # is dict, more than one key
+            raise ValueError("Spacy prototype data must dictionary with 1 key")
+        elif self.prototype_data is None:
             text_column_name = "text"
-
         else:
-            if (
-                isinstance(self.prototype_data, pd.DataFrame)
-                and len(self.prototype_data.columns) != 1
-            ):
-                raise TypeError("Expected 1 column of text data")
-
             text_column_name = (
                 self.prototype_data.columns[0]
                 if isinstance(self.prototype_data, pd.DataFrame)

vetiver/tests/test_spacy.py

@@ -32,7 +32,26 @@ def spacy_model():
     return nlp
 
 
-@pytest.fixture
+@pytest.fixture(scope="function")
+@pytest.mark.parametrize("data", ["a", 1, [1, 2, 3]])
+def test_bad_prototype_data(data, spacy_model):
+    with pytest.raises(TypeError):
+        vetiver.VetiverModel(spacy_model, "animals", prototype_data=data)
+
+
+@pytest.fixture(scope="function")
+@pytest.mark.parametrize(
+    pd.DataFrame(
+        {"col": ["1", "2"], "col2": [1, 2]},
+        pd.DataFrame({"col": ["1", "2"], "col2": [1, 2]}),
+    )
+)
+def test_bad_prototype_shape(data, spacy_model):
+    with pytest.raises(ValueError):
+        vetiver.VetiverModel(spacy_model, "animals", prototype_data=data)
+
+
+@pytest.fixture()
 def vetiver_client_with_prototype(spacy_model):  # With check_prototype=True
     df = pd.DataFrame({"new_column": ["one", "two", "three"]})
     v = vetiver.VetiverModel(spacy_model, "animals", prototype_data=df)
@@ -43,6 +62,16 @@ def vetiver_client_with_prototype(spacy_model):  # With check_prototype=True
     return client
 
 
+@pytest.fixture(scope="function")
+def vetiver_client_with_prototype_series(spacy_model):  # With check_prototype=True
+    df = pd.Series({"new_column": ["one", "two", "three"]})
+    v = vetiver.VetiverModel(spacy_model, "animals", prototype_data=df)
+    app = vetiver.VetiverAPI(v, check_prototype=True)
+    app.app.root_path = "/predict"
+    client = TestClient(app.app)
+    return client
+
+
 @pytest.fixture
 def vetiver_client_no_prototype(spacy_model):  # With check_prototype=False
     v = vetiver.VetiverModel(spacy_model, "animals")
@@ -80,6 +109,33 @@ def test_vetiver_predict_with_prototype(vetiver_client_with_prototype):
     }
 
 
+def test_vetiver_predict_with_prototype_series(vetiver_client_with_prototype_series):
+    df = pd.DataFrame({"new_column": ["turtles", "i have a dog"]})
+
+    response = vetiver.predict(endpoint=vetiver_client_with_prototype_series, data=df)
+
+    assert isinstance(response, pd.DataFrame), response
+    assert response.to_dict() == {
+        "0": {
+            "text": "turtles",
+            "ents": [],
+            "sents": [{"start": 0, "end": 7}],
+            "tokens": [{"id": 0, "start": 0, "end": 7}],
+        },
+        "1": {
+            "text": "i have a dog",
+            "ents": [{"start": 9, "end": 12, "label": "ANIMAL"}],
+            "sents": nan,
+            "tokens": [
+                {"id": 0, "start": 0, "end": 1},
+                {"id": 1, "start": 2, "end": 6},
+                {"id": 2, "start": 7, "end": 8},
+                {"id": 3, "start": 9, "end": 12},
+            ],
+        },
+    }
+
+
 def test_vetiver_predict_no_prototype(vetiver_client_no_prototype):
     df = pd.DataFrame({"uhhh": ["turtles", "i have a dog"]})