removing deprecated calls

zenml-io · Nov 13, 2024 · d399790 · d399790
1 parent 82f7c12
commit d399790
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 10 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -55,5 +55,5 @@ jobs:
         uses: ./.github/actions/llm_finetuning_template_test
         with:
           python-version: ${{ matrix.python-version }}
-          ref-zenml: ${{ inputs.ref-zenml || 'develop' }}
+          ref-zenml: ${{ inputs.ref-zenml || 'feature/followup-run-metadata' }}
           ref-template: ${{ inputs.ref-template || github.ref }}
diff --git a/template/steps/log_metadata.py b/template/steps/log_metadata.py
@@ -2,7 +2,7 @@
 
 from typing import Any, Dict
 
-from zenml import log_model_metadata, step, get_step_context
+from zenml import log_metadata, step, get_step_context
 
 
 @step(enable_cache=False)
@@ -24,4 +24,9 @@ def log_metadata_from_step_artifact(
 
     metadata = {artifact_name: metadata_dict}
 
-    log_model_metadata(metadata)
+    if context.model:
+        log_metadata(
+            metadata=metadata,
+            model_name=context.model.name,
+            model_version=context.model.version,
+        )
diff --git a/template/steps/prepare_datasets.py b/template/steps/prepare_datasets.py
@@ -6,7 +6,7 @@
 from materializers.directory_materializer import DirectoryMaterializer
 from typing_extensions import Annotated
 from utils.tokenizer import generate_and_tokenize_prompt, load_tokenizer
-from zenml import log_model_metadata, step
+from zenml import get_step_context, log_metadata, step
 from zenml.materializers import BuiltInMaterializer
 from zenml.utils.cuda_utils import cleanup_gpu_memory
 
@@ -33,12 +33,16 @@ def prepare_data(
 
     cleanup_gpu_memory(force=True)
 
-    log_model_metadata(
-        {
-            "system_prompt": system_prompt,
-            "base_model_id": base_model_id,
-        }
-    )
+    context = get_step_context()
+    if context.model:
+        log_metadata(
+            metadata={
+                "system_prompt": system_prompt,
+                "base_model_id": base_model_id,
+            },
+            model_name=context.model.name,
+            model_version=context.model.version,
+        )
 
     tokenizer = load_tokenizer(base_model_id, False, use_fast)
     gen_and_tokenize = partial(