Update model weight validation logic to handle special weight file naming (#13256)
This commit is contained in:
@@ -202,22 +202,33 @@ def validate_model_shards(model_path: Path) -> Tuple[bool, Optional[str], List[P
|
|||||||
)
|
)
|
||||||
|
|
||||||
if not shard_files:
|
if not shard_files:
|
||||||
# No sharded files - check for single model file
|
# No sharded files - check for any safetensors or bin files
|
||||||
single_files = list(model_path.glob("model.safetensors")) or list(
|
# Exclude non-model files like tokenizer, config, optimizer, etc.
|
||||||
model_path.glob("pytorch_model.bin")
|
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:
|
if single_files:
|
||||||
# Validate the single safetensors file if it exists
|
# Validate all safetensors files, not just the first one
|
||||||
if single_files[0].suffix == ".safetensors":
|
for model_file in single_files:
|
||||||
is_valid, error_msg = validate_safetensors_file(single_files[0])
|
if model_file.suffix == ".safetensors":
|
||||||
|
is_valid, error_msg = validate_safetensors_file(model_file)
|
||||||
if not is_valid:
|
if not is_valid:
|
||||||
return (
|
return (
|
||||||
False,
|
False,
|
||||||
f"Corrupted file {single_files[0].name}: {error_msg}",
|
f"Corrupted file {model_file.name}: {error_msg}",
|
||||||
[single_files[0]],
|
[model_file],
|
||||||
)
|
)
|
||||||
return True, None, []
|
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
|
# Extract total shard count from any shard filename
|
||||||
total_shards = None
|
total_shards = None
|
||||||
|
|||||||
Reference in New Issue
Block a user