def check_safetensors_precision(
    model_path: str,
    model: torch.nn.Module,
    expected_precision: torch.dtype,
    tolerance: float = 0.2,
):
    """Verify that the safetensors files in model_path have the expected size.

    Args:
        model_path: Path to the directory containing the safetensors files
        model: The original model to count parameters from
        expected_precision: The expected precision (torch.float32, torch.bfloat16, etc.)
        tolerance: Allowed deviation from expected file size (as a ratio)

    Returns:
        bool: True if the safetensors files have the expected size, False otherwise
    """
    total_params = sum(p.numel() for p in model.parameters())
    # Determine expected bytes per parameter based on precision
    bytes_per_param = {
        torch.float32: 4,
        torch.float16: 2,
        torch.bfloat16: 2,
        torch.int8: 1,
    }.get(expected_precision)
    assert bytes_per_param

    expected_size = total_params * bytes_per_param

    safetensors_files = glob.glob(os.path.join(model_path, '*.safetensors'))
    if not safetensors_files:
        # If no safetensors files found, check pytorch_model.bin
        safetensors_files = glob.glob(
            os.path.join(model_path, 'pytorch_model*.bin'),
        )

    if not safetensors_files:
        return False

    total_size = sum(os.path.getsize(f) for f in safetensors_files)
    size_ratio = total_size / expected_size

    is_correct_size = (1.0 - tolerance) <= size_ratio <= (1.0 + tolerance)
    return is_correct_size
Verify that the safetensors files in model_path have the expected size.

    Args:
        model_path: Path to the directory containing the safetensors files
        model: The original model to count parameters from
        expected_precision: The expected precision (torch.float32, torch.bfloat16, etc.)
        tolerance: Allowed deviation from expected file size (as a ratio)

    Returns:
        bool: True if the safetensors files have the expected size, False otherwise
    
tests/a_scripts/inference/test_convert_composer_to_hf.py