[Bug Fix] Download config.json for Large TimeSeries Model (#16295)
diff --git a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py b/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py index f1bf495..64975c2 100644 --- a/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py +++ b/iotdb-core/ainode/ainode/core/model/built_in_model_factory.py
@@ -52,12 +52,12 @@ logger = Logger() -def _download_file_from_hf_if_necessary( - file_path: str, local_dir: str, repo_id: str -) -> bool: - if not os.path.exists(file_path): +def _download_file_from_hf_if_necessary(local_dir: str, repo_id: str) -> bool: + weights_path = os.path.join(local_dir, MODEL_WEIGHTS_FILE_IN_SAFETENSORS) + config_path = os.path.join(local_dir, MODEL_CONFIG_FILE_IN_JSON) + if not os.path.exists(weights_path): logger.info( - f"Model file not found at {file_path}, downloading from HuggingFace..." + f"Model weights file not found at {weights_path}, downloading from HuggingFace..." ) try: hf_hub_download( @@ -65,9 +65,27 @@ filename=MODEL_WEIGHTS_FILE_IN_SAFETENSORS, local_dir=local_dir, ) - logger.info(f"Got file to {file_path}") + logger.info(f"Got file to {weights_path}") except Exception as e: - logger.error(f"Failed to download model file to {local_dir} due to {e}") + logger.error( + f"Failed to download model weights file to {local_dir} due to {e}" + ) + return False + if not os.path.exists(config_path): + logger.info( + f"Model config file not found at {config_path}, downloading from HuggingFace..." + ) + try: + hf_hub_download( + repo_id=repo_id, + filename=MODEL_CONFIG_FILE_IN_JSON, + local_dir=local_dir, + ) + logger.info(f"Got file to {config_path}") + except Exception as e: + logger.error( + f"Failed to download model config file to {local_dir} due to {e}" + ) return False return True @@ -82,11 +100,7 @@ bool: True if the model is existed or downloaded successfully, False otherwise. """ repo_id = TIMER_REPO_ID[model_type] - weights_path = os.path.join(local_dir, MODEL_WEIGHTS_FILE_IN_SAFETENSORS) - if not _download_file_from_hf_if_necessary(weights_path, local_dir, repo_id): - return False - config_path = os.path.join(local_dir, MODEL_CONFIG_FILE_IN_JSON) - if not _download_file_from_hf_if_necessary(config_path, local_dir, repo_id): + if not _download_file_from_hf_if_necessary(local_dir, repo_id): return False return True