RNN 人名分类案例 - 构建数据集对象

class NameClassDataset(Dataset):
    def __init__(self, my_list_x, my_list_y):
        ...
    def __len__(self):
        ...
    def __getitem__(self, index):
        ...

index = min(max(index, 0), self.sample_len - 1)

tensor_x = one_hot(name)  # 人名的 one-hot 编码
tensor_y = tensor(country)  # 国家的张量表示

import torch
from torch.utils.data import Dataset


class NameClassDataset(Dataset):
    """人名分类数据集"""
    
    def __init__(self, my_list_x, my_list_y, all_letters, all_countries):
        """
        初始化函数
        
        Args:
            my_list_x: 样本数据列表（人名）
            my_list_y: 标签数据列表（国家）
            all_letters: 全局字母表
            all_countries: 国家列表
        """
        # 1. 调用父类初始化
        super(NameClassDataset, self).__init__()
        
        # 2. 存储样本数据列表
        self.my_list_x = my_list_x
        
        # 3. 存储标签数据列表
        self.my_list_y = my_list_y
        
        # 4. 计算样本总数并存储
        self.sample_len = len(my_list_x)
        
        # 5. 存储字母表和国家列表
        self.all_letters = all_letters
        self.all_countries = all_countries
        self.n_letters = len(all_letters)
        self.n_countries = len(all_countries)
    
    def __len__(self):
        """获取样本总数"""
        return self.sample_len
    
    def __getitem__(self, index):
        """
        根据索引获取样本
        并进行 one-hot 编码和张量转换
        
        Args:
            index: 样本索引
        
        Returns:
            tensor_x: 人名的 one-hot 编码张量
            tensor_y: 国家的张量表示
        """
        # 1. 索引边界校验
        index = min(max(index, 0), self.sample_len - 1)
        
        # 2. 获取人名和国家
        name = self.my_list_x[index]
        country = self.my_list_y[index]
        
        # 3. one-hot 编码
        tensor_x = torch.zeros(len(name), self.n_letters)
        for i, letter in enumerate(name):
            tensor_x[i][self.all_letters.find(letter)] = 1
        
        # 4. 国家转索引
        tensor_y = torch.tensor(self.all_countries.index(country))
        
        return tensor_x, tensor_y


# 测试
if __name__ == "__main__":
    print("=== RNN 人名分类案例 - 构建数据集对象 ===")
    
    import string
    
    all_letters = string.ascii_letters + " .,;'"
    all_countries = ["Chinese", "English", "German", "French"]
    
    my_list_x = ["Zhang", "Smith", "Muller", "Martin"]
    my_list_y = ["Chinese", "English", "German", "French"]
    
    dataset = NameClassDataset(my_list_x, my_list_y, all_letters, all_countries)
    
    print(f"样本总数: {len(dataset)}")
    
    tensor_x, tensor_y = dataset[0]
    print(f"人名张量形状: {tensor_x.shape}")
    print(f"国家索引: {tensor_y}")