Merge pull request #2 from kristeligt-dagblad/new_models

Added support for AutoML, DNN and Boosted Tree models
kristeligt-dagblad · Jun 25, 2020 · fc119ff · fc119ff
2 parents 0038829 + d98f4d8
commit fc119ff
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 30 deletions.
diff --git a/dbt_project.yml b/dbt_project.yml
@@ -1,4 +1,6 @@
 name: "dbt_ml"
-version: "0.1.0"
+version: "0.3.0"
 
-require-dbt-version: [">=0.16.0"]
+config-version: 2
+
+require-dbt-version: [">=0.17.0"]
diff --git a/macros/hooks/model_audit.sql b/macros/hooks/model_audit.sql
@@ -13,33 +13,54 @@
 
 {% macro _audit_insert_templates() %}
 
-{% do return ({
-    'default': {
-        'training_info': [
-            'training_run',
-            'iteration',
-            'loss',
-            'eval_loss',
-            'learning_rate',
-            'duration_ms',
-            'array(select as struct null as centroid_id, cast(null as float64) as cluster_radius, null as cluster_size)'
-        ],
-        'feature_info': ['*'],
-        'weights': ['*']
-    },
-    'kmeans': {
-        'training_info': [
-            'training_run',
-            'iteration',
-            'loss',
-            'cast(null as float64) as eval_loss',
-            'cast(null as float64) as learning_rate',
-            'duration_ms',
-            'cluster_info'
-        ],
-        'feature_info': ['*']
-    }
-}) %}
+{%- set schemas -%}
+
+default: &default
+    training_info: &default_training_info
+        - training_run
+        - iteration
+        - loss
+        - eval_loss
+        - learning_rate
+        - duration_ms
+        - array(select as struct null as centroid_id, cast(null as float64) as cluster_radius, null as cluster_size)
+    feature_info: &default_feature_info ['*']
+automl_classifier:
+    training_info: *default_training_info
+    feature_info: *default_feature_info
+    weights: ['*']
+automl_regressor:
+    training_info: *default_training_info
+    feature_info: *default_feature_info
+    weights: ['*']
+boosted_tree_classifier: *default
+boosted_tree_regressor: *default
+dnn_classifier: *default
+dnn_regressor: *default
+kmeans:
+    training_info:
+        - training_run
+        - iteration
+        - loss
+        - eval_loss
+        - learning_rate
+        - duration_ms
+        - cluster_info
+    feature_info: *default_feature_info
+linear_reg:
+    training_info: *default_training_info
+    feature_info: *default_feature_info
+    weights: ['*']
+logistic_reg:
+    training_info: *default_training_info
+    feature_info: *default_feature_info
+    weights: ['*']
+matrix_factorization: *default
+tensorflow: {}
+
+{%- endset -%}
+
+{% do return(fromyaml(schemas)) %}
 
 {% endmacro %}
 
@@ -63,7 +84,7 @@
 {% macro model_audit() %}
 
     {% set model_type = config.get('ml_config')['model_type'] %}
-    {% set model_type_repr = dbt_ml._audit_insert_templates().get(model_type, 'default') %}
+    {% set model_type_repr = model_type if model_type in dbt_ml._audit_insert_templates().keys() else 'default' %}
 
     {% set info_types = ['training_info', 'feature_info', 'weights'] %}