全球人名分类案例 - RNN 模型预测

import torch


def name_to_tensor(name, all_letters, n_letters):
    """
    将人名转换为 one-hot 张量
    
    Args:
        name: 人名
        all_letters: 全局字母表
        n_letters: 字母表大小
    
    Returns:
        tensor: one-hot 张量
    """
    tensor = torch.zeros(len(name), 1, n_letters)
    for i, letter in enumerate(name):
        tensor[i][0][all_letters.find(letter)] = 1
    return tensor


def my_predict(name, model, all_letters, n_letters, all_countries, top_k=3):
    """
    预测人名来自哪个国家
    
    Args:
        name: 人名
        model: 训练好的模型
        all_letters: 全局字母表
        n_letters: 字母表大小
        all_countries: 国家列表
        top_k: 返回前 k 名
    
    Returns:
        前 k 名结果
    """
    # 1. 将人名转换为 one-hot 张量
    input_tensor = name_to_tensor(name, all_letters, n_letters)
    
    # 2. 加载模型参数
    model.load_state_dict(torch.load('model_params.pth'))
    model.eval()
    
    # 3. 预测
    with torch.no_grad():
        output = model(input_tensor)
        
        # 4. 返回前 k 名
        topv, topi = output.topk(top_k, dim=1, largest=True)
        
        for i in range(top_k):
            value = topv[0][i]
            index = topi[0][i]
            country = all_countries[index]
            print(f"第 {i+1} 名: {country} (概率: {value:.4f})")


# 测试
if __name__ == "__main__":
    print("=== 全球人名分类案例 - RNN 模型预测 ===")
    
    import string
    
    all_letters = string.ascii_letters + " .,;'"
    n_letters = len(all_letters)
    all_countries = ["Chinese", "English", "German", "French"]
    
    # 假设 model 已定义
    # model = RNNModel(n_letters, 64, n_countries)
    
    # 预测
    # my_predict("Zhang", model, all_letters, n_letters, all_countries, top_k=3)
    # my_predict("Smith", model, all_letters, n_letters, all_countries, top_k=3)
    # my_predict("Muller", model, all_letters, n_letters, all_countries, top_k=3)