[#60] Hook up the new save model metadata training step 3

ipums · Oct 18, 2023 · 8622b44 · 8622b44
1 parent d22024f
commit 8622b44
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/hlink/linking/table_definitions.csv b/hlink/linking/table_definitions.csv
@@ -18,6 +18,7 @@ training_data,Training: Raw training data as read from specified file,1
 hh_training_data,Household Training: Raw HH training data as read from specified file,1
 training_features,Training: Transformed data with features,1
 hh_training_features,Household Training: Transformed HH training data with features,1
+training_model_metadata,Training: Metadata like feature weights that describe the trained model,0
 model_eval_training_data,Model Exploration: Raw training data as read from specified file,1
 model_eval_training_vectorized,Model Exploration: Training data after applying comparison feature and pipeline transformations,1
 model_eval_training_results,Model Exploration: Results of ML model exploration train/test splits,0

diff --git a/hlink/linking/training/link_step_save_model_metadata.py b/hlink/linking/training/link_step_save_model_metadata.py
@@ -13,11 +13,8 @@ def __init__(self, task):
         super().__init__(
             task,
             "save metadata about the model",
-            input_table_names=[
-                f"{task.table_prefix}training_features",
-                f"{task.table_prefix}training_vectorized",
-            ],
-            output_table_names=[f"{task.table_prefix}training_results"],
+            output_table_names=[f"{task.table_prefix}training_model_metadata"],
+            input_model_names=[f"{task.table_prefix}trained_model"],
         )
 
     def _run(self):

diff --git a/hlink/linking/training/training.py b/hlink/linking/training/training.py
@@ -8,6 +8,7 @@
 from .link_step_ingest_file import LinkStepIngestFile
 from .link_step_create_comparison_features import LinkStepCreateComparisonFeatures
 from .link_step_train_and_save_model import LinkStepTrainAndSaveModel
+from .link_step_save_model_metadata import LinkStepSaveModelMetadata
 
 
 class Training(LinkTask):
@@ -21,4 +22,5 @@ def get_steps(self):
             LinkStepIngestFile(self),
             LinkStepCreateComparisonFeatures(self),
             LinkStepTrainAndSaveModel(self),
+            LinkStepSaveModelMetadata(self),
         ]