diff --git a/scripts/ci/validate_and_download_models.py b/scripts/ci/validate_and_download_models.py index 31f035e77..91678615a 100755 --- a/scripts/ci/validate_and_download_models.py +++ b/scripts/ci/validate_and_download_models.py @@ -202,22 +202,33 @@ def validate_model_shards(model_path: Path) -> Tuple[bool, Optional[str], List[P ) if not shard_files: - # No sharded files - check for single model file - single_files = list(model_path.glob("model.safetensors")) or list( - model_path.glob("pytorch_model.bin") - ) + # No sharded files - check for any safetensors or bin files + # Exclude non-model files like tokenizer, config, optimizer, etc. + all_safetensors = list(model_path.glob("*.safetensors")) + all_bins = list(model_path.glob("*.bin")) + + # Filter out non-model files + excluded_prefixes = ["tokenizer", "optimizer", "training_", "config"] + single_files = [ + f + for f in (all_safetensors or all_bins) + if not any(f.name.startswith(prefix) for prefix in excluded_prefixes) + and not f.name.endswith(".index.json") + ] + if single_files: - # Validate the single safetensors file if it exists - if single_files[0].suffix == ".safetensors": - is_valid, error_msg = validate_safetensors_file(single_files[0]) - if not is_valid: - return ( - False, - f"Corrupted file {single_files[0].name}: {error_msg}", - [single_files[0]], - ) + # Validate all safetensors files, not just the first one + for model_file in single_files: + if model_file.suffix == ".safetensors": + is_valid, error_msg = validate_safetensors_file(model_file) + if not is_valid: + return ( + False, + f"Corrupted file {model_file.name}: {error_msg}", + [model_file], + ) return True, None, [] - return False, "No model files found (safetensors or bin)", [] + return False, "No model weight files found (safetensors or bin)", [] # Extract total shard count from any shard filename total_shards = None