Fix import and mocking (#1169)

mosaicml · May 3, 2024 · c0d591c · c0d591c
1 parent ddf4aa4
commit c0d591c
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/llmfoundry/eval/datasets/in_context_learning_evaluation.py b/llmfoundry/eval/datasets/in_context_learning_evaluation.py
@@ -15,7 +15,6 @@
 import transformers
 from composer.core import DataSpec
 from composer.core.data_spec import _default_split_batch, _split_list
-from composer.datasets.utils import stop_sequences_criteria
 from composer.utils import MissingConditionalImportError, dist, get_file
 from datasets import Dataset as HFDataset
 from datasets import IterableDataset, load_dataset
@@ -26,6 +25,7 @@
     get_continuation_span,
     get_fewshot_sample_idxs,
     make_padded_input,
+    stop_sequences_criteria,
     strip_data,
     tokenizer_needs_prefix_space,
     trim_context,

diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py
@@ -875,6 +875,8 @@ def test_huggingface_conversion_callback(
     mlflow_logger_mock.model_registry_prefix = ''
     mlflow_logger_mock._experiment_id = 'mlflow-experiment-id'
     mlflow_logger_mock._run_id = 'mlflow-run-id'
+    mlflow_logger_mock._enabled = True
+    mlflow_logger_mock.run_url = 'fake-url'
     trainer = Trainer(
         model=original_model,
         device='gpu',