FineTurnX/test_quantization.py

"""
测试量化模块
"""

import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from finetunex.quantization import (
    get_model_size,
    estimate_quantized_size,
    compare_models,
)


def test_get_model_size():
    """测试获取模型大小"""
    print("=" * 60)
    print("测试：get_model_size")
    print("=" * 60)
    
    # 使用示例数据目录作为测试
    test_path = "./data"
    
    if os.path.exists(test_path):
        size_info = get_model_size(test_path)
        print(f"路径：{test_path}")
        print(f"总大小：{size_info['total_size_formatted']}")
        print(f"文件数：{size_info['file_count']}")
        print("✓ 测试通过\n")
    else:
        print(f"⚠ 测试路径不存在：{test_path}\n")


def test_estimate_quantized_size():
    """测试估算量化后大小"""
    print("=" * 60)
    print("测试：estimate_quantized_size")
    print("=" * 60)
    
    test_path = "./data"
    
    if os.path.exists(test_path):
        print(f"路径：{test_path}")
        
        for bits in [4, 8]:
            estimate = estimate_quantized_size(test_path, quantization_bits=bits)
            print(f"\n{bits}bit 量化估算:")
            print(f"  原始大小：{estimate['original_size']}")
            print(f"  估算大小：{estimate['estimated_size']}")
            print(f"  压缩比：{estimate['compression_ratio']}")
            print(f"  节省空间：{estimate['space_saved']} ({estimate['space_saved_percent']})")
        
        print("\n✓ 测试通过\n")
    else:
        print(f"⚠ 测试路径不存在：{test_path}\n")


def test_compare_models():
    """测试比较模型大小"""
    print("=" * 60)
    print("测试：compare_models")
    print("=" * 60)
    
    # 比较两个目录
    path1 = "./data"
    path2 = "./configs"
    
    if os.path.exists(path1) and os.path.exists(path2):
        comparison = compare_models(path1, path2, "数据目录", "配置目录")
        
        print(f"数据目录：{comparison['数据目录']['size']}")
        print(f"配置目录：{comparison['配置目录']['size']}")
        print(f"差异：{comparison['difference']} ({comparison['difference_percent']})")
        print(f"更小：{comparison['smaller']}")
        print("\n✓ 测试通过\n")
    else:
        print(f"⚠ 测试路径不存在\n")


def main():
    print("\n" + "=" * 60)
    print("量化模块测试")
    print("=" * 60 + "\n")
    
    # 运行测试
    test_get_model_size()
    test_estimate_quantized_size()
    test_compare_models()
    
    print("=" * 60)
    print("所有测试完成！")
    print("=" * 60)
    print("\n提示：这些测试使用了示例目录，实际使用时请指定模型路径。")
    print("例如：--model_path ./outputs/qwen3.5-0.8b-finetuned\n")


if __name__ == "__main__":
    main()