Train extract v3

tests/connectors/test_train.py

@@ -154,6 +154,27 @@ def test_extract_write_backwards_compatible(self):
         assert df.iloc[0]['Entity to Find'] == 'Rachel'
 
     def test_extract_write(self):
+        """
+        Writing data to a wrangle (re-training)
+        """
+        recipe = """
+        write:
+        - train.extract:
+            columns:
+                - Entity to Find
+                - Variation (Optional)
+                - Notes
+            model_id: ee5f020e-d88e-4bd5
+        """
+        data = pd.DataFrame({
+            'Entity to Find': ['Rachel', 'Dolores', 'TARS'],
+            'Variation (Optional)': ['', '', ''],
+            'Notes': ['Blade Runner', 'Westworld', 'Interstellar'],
+        })
+        df = wrangles.recipe.run(recipe, dataframe=data)
+        assert df.iloc[0]['Entity to Find'] == 'Rachel'
+
+    def test_extract_write_2_0(self):
         """
         Writing data to a wrangle (re-training)
         """
@@ -174,6 +195,27 @@ def test_extract_write(self):
         df = wrangles.recipe.run(recipe, dataframe=data)
         assert df.iloc[0]['Find'] == 'Rachel'
 
+    def test_extract_write_3_0(self):
+        """
+        Writing data to a wrangle (re-training)
+        """
+        recipe = """
+        write:
+        - train.extract:
+            columns:
+                - Find
+                - Output
+                - Notes
+            model_id: ee5f020e-d88e-4bd5
+        """
+        data = pd.DataFrame({
+            'Find': ['Rachel', 'Dolores', 'TARS'],
+            'Output': ['', '', ''],
+            'Notes': ['Blade Runner', 'Westworld', 'Interstellar'],
+        })
+        df = wrangles.recipe.run(recipe, dataframe=data)
+        assert df.iloc[0]['Find'] == 'Rachel'
+
     def test_extract_ai_write(self):
         """
         Writing data to a wrangle (re-training)

wrangles/connectors/train.py

@@ -169,6 +169,7 @@ def write(
 
         if variant == 'pattern':
             versions = [
+                {'columns': ['Find', 'Output', 'Notes'], 'version': 'pattern 3.0'},
                 {'columns': ['Find', 'Output (Optional)', 'Notes'], 'version': 'pattern 2.0'},
                 {'columns': ['Entity to Find', 'Variation (Optional)', 'Notes'], 'version': 'pattern 1.0'}
             ]

wrangles/train.py

@@ -71,8 +71,8 @@ def extract(training_data: list, name: str = None, model_id: str = None, variant
                 raise ValueError(f"Training_data list must contain a list of two elements, plus optional Notes. Check element(s) {check_index} in training_list.\nFormat:\nFirst element is 'Entity to Find'\nSecond Element is 'Variation', If no variation, use \'\'\n"
                 "Example:[['Television', 'TV', '']]")
             # checking the first element in training list
-            if training_data[0] != ['Find', 'Output (Optional)', 'Notes']:
-                training_data = [['Find', 'Output (Optional)', 'Notes']] + training_data
+            if training_data[0] != ['Find', 'Output', 'Notes']:
+                training_data = [['Find', 'Output', 'Notes']] + training_data
 
         # If input is a list, check to make sure that all sublists are length of 7
         # Must have all values filled ('' counts as filled, None does not count)