fix: get_model in ray executor

modelscope · Dec 4, 2023 · ef34418 · ef34418
1 parent d0f564f
commit ef34418
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/data_juicer/utils/model_utils.py b/data_juicer/utils/model_utils.py
@@ -285,4 +285,8 @@ def get_model(model_key, lang='en', model_type='sentencepiece'):
 
     :param model_key: name of the model or tokenzier
     """
+    if model_key is None:
+        return None
+    if model_key not in MODEL_ZOO:
+        prepare_model(lang=lang, model_type=model_type, model_key=model_key)
     return MODEL_ZOO.get(model_key, None)