test_llama_model.py

#!/usr/bin/env python3
"""
测试Llama3.1-8B模型加载和基本功能
"""

import sys
import os
sys.path.append('/home/czy/CISC_PRO')

from src.models import get_model, InferenceConfig

def test_llama_model():
    """测试Llama3.1-8B模型"""
    print("🚀 开始测试Llama3.1-8B模型...")
    
    try:
        # 加载模型
        print("📂 正在加载模型...")
        model = get_model("llama3.1-8b")
        print("✅ 模型加载成功!")
        
        # 配置推理参数
        config = InferenceConfig(
            num_samples=1,
            temperature=0.7,
            top_p=0.9,
            max_new_tokens=256
        )
        
        # 测试简单的数学问题
        test_question = "Sarah has 10 apples. She gives 3 to her friend and buys 5 more. How many apples does she have now?"
        
        print(f"🧮 测试问题: {test_question}")
        print("💭 正在生成回答...")
        
        response = model.generate_response(f"Question: {test_question}\n\nAnswer:", config)
        
        print("✅ 回答生成成功!")
        print(f"🎯 模型回答: {response}")
        
        # 测试token计数
        token_count = model.count_tokens(test_question)
        print(f"📊 问题token数: {token_count}")
        
        print("\n🎉 Llama3.1-8B模型测试通过!")
        return True
        
    except Exception as e:
        print(f"❌ 模型测试失败: {str(e)}")
        return False

if __name__ == "__main__":
    success = test_llama_model()
    if success:
        print("\n✅ 您可以开始使用Llama3.1-8B进行路由系统评估!")
    else:
        print("\n❌ 请检查模型配置和路径设置")