Train extract v3

tests/connectors/test_train.py

@@ -188,6 +188,27 @@ def test_extract_write_backwards_compatible_adding_optional_label(self):
 
 
     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)
         """

wrangles/connectors/train.py

@@ -102,9 +102,9 @@ def read(model_id: str) -> _pd.DataFrame:
         elif len(tmp_data['Data'][0]) == 2:
             # Add a third column for Notes of empty strings
             [x.append('') for x in tmp_data['Data']]
-            columns = ['Entity to Find', 'Variation (Optional)', 'Notes']
+            columns = ['Find', 'Output', 'Notes']
         elif len(tmp_data['Data'][0]) == 3:
-            columns = ['Entity to Find', 'Variation (Optional)', 'Notes']
+            columns = ['Find', 'Output', 'Notes']
         else:
             raise ValueError("Extract Wrangle data should contain three columns. Check Wrangle data")
 
@@ -169,17 +169,17 @@ def write(
 
         if variant == 'pattern':
             versions = [
-                {'columns': ['Find', 'Output (Optional)', 'Notes'], 'version': 'pattern 2.0'},
-                {'columns': ['Entity to Find', 'Variation (Optional)', 'Notes'], 'version': 'pattern 1.0'},
                 {'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'}
             ]
             try:
                 required_columns = [
                         version for version in versions
                         if set(version["columns"]).issubset(set(df.columns.to_list()))
                     ][0]['columns']
             except:
-                required_columns = ['Find', 'Output (Optional)', 'Notes']
+                required_columns = ['Find', 'Output', 'Notes']
             col_len = 3
         elif variant == 'extract-ai':
             required_columns = ['Find', 'Description', 'Type', 'Default', 'Examples', 'Enum', 'Notes']

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)