如何使用torchtext v0.9实现自定义数据集？基于SMS垃圾短信数据集场景

TorchText v0.9版本重构了核心API，废弃了原有v0.8及之前的Field、TabularDataset等接口，以下是基于新API适配SMS Spam Collection数据集的完整实现流程：

前置依赖 #

• PyTorch >= 1.8.0
• TorchText == 0.9.0

实现步骤 #

1. 数据集格式说明 #

SMS Spam Collection数据集每行由制表符分隔为两部分：第一列为标签，取值为ham（正常短信）或spam（垃圾短信），第二列为短信原始文本内容。

2. 完整代码实现 #

import torch
from torch.utils.data import Dataset, DataLoader
from torchtext.vocab import build_vocab_from_iterator
from torchtext.transforms import VocabTransform, ToTensor

# 1. 加载原始数据集
def load_sms_data(file_path):
    labels = []
    texts = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            label, text = line.split('\t', 1)
            labels.append(label)
            texts.append(text)
    return labels, texts

# 2. 分词函数（可替换为spaCy等更专业的分词工具）
def tokenizer(text):
    return text.lower().split()

# 3. 构建自定义Dataset类
class SMSDataset(Dataset):
    def __init__(self, labels, texts, vocab=None, max_vocab_size=10000):
        self.labels = labels
        self.texts = texts
        # 标签映射规则
        self.label2id = {'ham': 0, 'spam': 1}
        # 构建词汇表：训练集传入vocab=None生成新词汇表，验证/测试集复用训练集词汇表
        if vocab is None:
            self.vocab = build_vocab_from_iterator(
                map(tokenizer, self.texts),
                max_tokens=max_vocab_size,
                specials=['<unk>', '<pad>']
            )
            self.vocab.set_default_index(self.vocab['<unk>'])
        else:
            self.vocab = vocab
        # 文本转换流程：转id -> 转tensor
        self.text_transform = torch.nn.Sequential(
            VocabTransform(self.vocab),
            ToTensor(dtype=torch.long)
        )
    
    def __len__(self):
        return len(self.labels)
    
    def __getitem__(self, idx):
        text = self.texts[idx]
        label = self.label2id[self.labels[idx]]
        tokenized = tokenizer(text)
        text_ids = self.text_transform(tokenized)
        return text_ids, torch.tensor(label, dtype=torch.long)

# 4. 批次对齐函数：解决不同长度文本的padding问题
def collate_batch(batch):
    text_list, label_list = [], []
    for text_ids, label in batch:
        text_list.append(text_ids)
        label_list.append(label)
    # 按批次内最大长度padding，<pad>对应的id为1
    padded_texts = torch.nn.utils.rnn.pad_sequence(text_list, batch_first=True, padding_value=1)
    return padded_texts, torch.stack(label_list)

# 使用示例
if __name__ == "__main__":
    # 替换为你本地的SMSSpamCollection文件路径
    labels, texts = load_sms_data('./SMSSpamCollection')
    # 8:2划分训练集和测试集
    train_size = int(0.8 * len(labels))
    train_labels, test_labels = labels[:train_size], labels[train_size:]
    train_texts, test_texts = texts[:train_size], texts[train_size:]
    # 构造数据集
    train_dataset = SMSDataset(train_labels, train_texts)
    test_dataset = SMSDataset(test_labels, test_texts, vocab=train_dataset.vocab)
    # 构造DataLoader
    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, collate_fn=collate_batch)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, collate_fn=collate_batch)
    
    # 测试输出
    for batch_text, batch_label in train_loader:
        print("批次文本shape:", batch_text.shape)
        print("批次标签shape:", batch_label.shape)
        break

注意事项 #

• 训练集生成的vocab必须复用到验证集、测试集，避免出现词汇表不一致的问题
• 如果需要文本截断逻辑，可以在__getitem__方法中对转id后的序列做切片处理
• 中文场景下可将分词函数替换为jieba等中文分词工具

内容的提问来源于stack exchange，提问作者Charitarth Chugh