前言
由于计算机无法认识到文字内容,因此在训练模型时需要将文字映射到计算机能够识别的编码内容。
映射的流程如下:
- 首先将文字内容按照词表映射到成唯一的数字ID。比如"我爱中国",将"中"映射为1,将"国"映射到2。
- 再将文字映射到的数字ID映射成向量。比如"中"映射成了ID=1,再将1映射成某个向量,比如[0,0,0,0,0,1]。
python
# 构建数据集:词汇表、训练集、验证集、测试集
vocab, train_data, dev_data, test_data = build_dataset(config, args.word)
# 构建数据迭代器,用于批量加载数据
train_iter = build_iterator(train_data, config)
dev_iter = build_iterator(dev_data, config)
test_iter = build_iterator(test_data, config)
这篇文章的目的就是搞懂上面的代码如何实现,即如何构建文字数据集和迭代器。
源码
python
# coding: UTF-8
# coding: UTF-8
# 导入必要的库
import os # 操作系统接口,用于文件路径处理
import torch # PyTorch深度学习框架
import numpy as np # 数值计算库
import pickle as pkl # 对象序列化/反序列化,用于保存词汇表
from tqdm import tqdm # 进度条显示
import time
from datetime import timedelta
# 全局常量定义
MAX_VOCAB_SIZE = 10000 # 词汇表最大容量限制
UNK, PAD = '<UNK>', '<PAD>' # 特殊标记:未知词(UNK)和填充符(PAD)
def build_vocab(file_path, tokenizer, max_size, min_freq):
"""构建词汇表字典
Args:
file_path: 训练集文件路径
tokenizer: 分词函数(按词或字符分割)
max_size: 最大词汇表大小
min_freq: 词的最小出现频次阈值
Returns:
vocab_dic: 词到索引的映射字典,包含UNK和PAD
"""
vocab_dic = {}
# 遍历训练集文件的每一行
with open(file_path, 'r', encoding='UTF-8') as f:
for line in tqdm(f): # 使用tqdm显示进度条
lin = line.strip()
if not lin:
continue # 跳过空行
content = lin.split('\t') # 分割文本和标签,取文本内容
# 分词并统计词频
for word in tokenizer(content):
vocab_dic[word] = vocab_dic.get(word, 0) + 1
# 筛选词频≥min_freq的词,按词频降序排列,取前max_size个
vocab_list = sorted([item for item in vocab_dic.items() if item >= min_freq],
key=lambda x: x, reverse=True)[:max_size]
# 生成词到索引的映射字典
vocab_dic = {word_count: idx for idx, word_count in enumerate(vocab_list)}
# 添加未知词和填充符的索引(排在最后两位)
vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})
return vocab_dic
def build_dataset(config, use_word):
"""构建数据集
Args:
config: 配置对象,包含文件路径等参数
use_word: 分词方式,True表示按词分割,False按字符分割
Returns:
vocab: 词汇表字典
train/dev/test: 处理后的数据集
"""
# 定义分词器
if use_word:
tokenizer = lambda x: x.split(' ') # 按空格分割(词级别)
else:
tokenizer = lambda x: [y for y in x] # 按字符分割
# 加载或创建词汇表
if os.path.exists(config.vocab_path):
vocab = pkl.load(open(config.vocab_path, 'rb')) # 从文件加载
else:
vocab = build_vocab(config.train_path, tokenizer, MAX_VOCAB_SIZE, min_freq=1)
pkl.dump(vocab, open(config.vocab_path, 'wb')) # 保存词汇表
print(f"Vocab size: {len(vocab)}")
# -----------------------n-gram哈希函数定义-----------------------
def biGramHash(sequence, t, buckets):
"""计算第t个位置的bigram哈希值
公式: (前一个词的哈希值 * 质数) % 桶大小
说明:如果t-1越界,用0代替(相当于用PAD的哈希值)
"""
t1 = sequence[t - 1] if t - 1 >= 0 else 0
return (t1 * 14918087) % buckets # 14918087是一个大质数
def triGramHash(sequence, t, buckets):
"""计算第t个位置的trigram哈希值
公式: (前两个词的哈希值组合 * 质数) % 桶大小
说明:如果t-1或t-2越界,用0代替
"""
t1 = sequence[t - 1] if t - 1 >= 0 else 0
t2 = sequence[t - 2] if t - 2 >= 0 else 0
return (t2 * 14918087 * 18408749 + t1 * 14918087) % buckets # 双质数减少冲突
# -------------------------------------------------------------
def load_dataset(path, pad_size=32):
"""加载并处理单个数据集文件
Args:
path: 数据集文件路径
pad_size: 填充/截断后的固定长度
Returns:
contents: 处理后的数据列表,元素为(词索引, 标签, 长度, bigram, trigram)
"""
contents = []
with open(path, 'r', encoding='UTF-8') as f:
for line in tqdm(f): # 显示进度条
lin = line.strip()
if not lin:
continue
content, label = lin.split('\t') # 分割文本和标签
# 分词并处理长度
token = tokenizer(content)
seq_len = len(token)
# 填充或截断至固定长度pad_size
if pad_size:
if len(token) < pad_size:
# 注意:这里用vocab.get(PAD)可能存在错误,PAD应为词汇表中已存在的键
token.extend([vocab.get(PAD)] * (pad_size - len(token))) # 填充
else:
token = token[:pad_size] # 截断
seq_len = pad_size # 更新实际长度为pad_size
# 将词转换为索引,未知词用UNK的索引
words_line = []
for word in token:
words_line.append(vocab.get(word, vocab.get(UNK))) # 双重保险取UNK索引
# 生成n-gram特征(FastText模型需要)
buckets = config.n_gram_vocab # 从配置获取哈希桶数量
bigram = []
trigram = []
for i in range(pad_size):
# 为每个位置生成bigram和trigram的哈希值
bigram.append(biGramHash(words_line, i, buckets))
trigram.append(triGramHash(words_line, i, buckets))
# 添加处理后的数据:词索引、标签、长度、bigram、trigram
contents.append((words_line, int(label), seq_len, bigram, trigram))
return contents # 返回结构:[([...], 0, 32, [...], [...]), ...]
# 加载并处理所有数据集
train = load_dataset(config.train_path, config.pad_size)
dev = load_dataset(config.dev_path, config.pad_size)
test = load_dataset(config.test_path, config.pad_size)
return vocab, train, dev, test
class DatasetIterater(object):
"""数据集迭代器,用于按批次生成数据"""
def __init__(self, batches, batch_size, device):
"""
Args:
batches: 处理后的数据集,格式为[(words_line, label, seq_len, bigram, trigram), ...]
batch_size: 每个批次的样本数
device: 数据存放设备(cpu或cuda)
"""
self.batch_size = batch_size
self.batches = batches
self.n_batches = len(batches) // batch_size # 完整批次数
self.residue = False # 是否包含不完整的剩余批次
# 如果总样本数不能被batch_size整除,设置residue标志
if len(batches) % self.n_batches != 0:
self.residue = True
self.index = 0 # 当前批次索引
self.device = device # 设备类型
def _to_tensor(self, datas):
"""将原始数据转换为Tensor格式"""
# 注释掉的代码为按序列长度排序的逻辑(可用于动态padding优化)
# xx = [xxx for xxx in datas] # 获取所有样本的原始长度
# indexx = np.argsort(xx)[::-1] # 按长度降序排列的索引
# datas = np.array(datas)[indexx] # 重新排列数据
# 构造各特征张量(LongTensor用于整型数据)
x = torch.LongTensor([_ for _ in datas]).to(self.device) # 词索引序列
y = torch.LongTensor([_ for _ in datas]).to(self.device) # 标签
bigram = torch.LongTensor([_ for _ in datas]).to(self.device) # bigram特征
trigram = torch.LongTensor([_ for _ in datas]).to(self.device) # trigram特征
# 实际长度(考虑padding前的原始长度,但不超过pad_size)
seq_len = torch.LongTensor([_ for _ in datas]).to(self.device)
return (x, seq_len, bigram, trigram), y # 返回特征元组和标签
def __next__(self):
"""生成下一个批次数据"""
# 处理剩余的不完整批次(当总样本数不是batch_size整数倍时)
if self.residue and self.index == self.n_batches:
batches = self.batches[self.index * self.batch_size: len(self.batches)]
self.index += 1
batches = self._to_tensor(batches)
return batches
# 所有批次处理完成后重置索引并抛出停止迭代异常
elif self.index >= self.n_batches:
self.index = 0
raise StopIteration
# 正常批次处理
else:
batches = self.batches[self.index * self.batch_size: (self.index + 1) * self.batch_size]
self.index += 1
batches = self._to_tensor(batches)
return batches
def __iter__(self):
"""返回迭代器自身"""
return self
def __len__(self):
"""返回总批次数(包含剩余批次)"""
return self.n_batches + 1 if self.residue else self.n_batches
def build_iterator(dataset, config):
"""构建数据集迭代器
Args:
dataset: 处理后的数据集
config: 配置对象,需包含batch_size和device属性
Returns:
DatasetIterater实例
"""
iter = DatasetIterater(dataset, config.batch_size, config.device)
return iter
def get_time_dif(start_time):
"""计算时间间隔
Args:
start_time: 开始时间戳
Returns:
timedelta: 格式化的时间差(秒级精度)
示例:
>>> start = time.time()
>>> # 执行操作...
>>> print(get_time_dif(start)) # 输出: 0:00:12
"""
end_time = time.time()
time_dif = end_time - start_time
return timedelta(seconds=int(round(time_dif)))
if __name__ == "__main__":
"""预训练词向量提取(示例用法)"""
# 文件路径配置
vocab_dir = "./THUCNews/data/vocab.pkl" # 词汇表路径
pretrain_dir = "./THUCNews/data/sgns.sogou.char" # 预训练向量路径
filename_trimmed_dir = "./THUCNews/data/vocab.embedding.sougou" # 输出路径
emb_dim = 300 # 词向量维度
# 加载词汇表(词到id的映射字典)
word_to_id = pkl.load(open(vocab_dir, 'rb'))
# 初始化随机词向量矩阵(词汇表大小 x 维度)
embeddings = np.random.rand(len(word_to_id), emb_dim)
# 加载预训练词向量
with open(pretrain_dir, "r", encoding='UTF-8') as f:
for i, line in enumerate(f.readlines()):
# 跳过首行标题(如果存在)
# if i == 0: continue
lin = line.strip().split(" ")
word = lin # 词
vector = lin[1:301] # 对应向量
# 如果当前词在词汇表中,更新其向量
if word in word_to_id:
idx = word_to_id[word]
emb = [float(x) for x in vector] # 转换为浮点数列表
embeddings[idx] = np.asarray(emb, dtype='float32') # 更新矩阵
# 保存压缩后的词向量矩阵(npz格式)
np.savez_compressed(filename_trimmed_dir, embeddings=embeddings)
数据集构建
python
def build_dataset(config, use_word):
"""构建数据集
Args:
config: 配置对象,包含文件路径等参数
use_word: 分词方式,True表示按词分割,False按字符分割
Returns:
vocab: 词汇表字典
train/dev/test: 处理后的数据集
"""
# 定义分词器
if use_word:
tokenizer = lambda x: x.split(' ') # 按空格分割(词级别)
else:
tokenizer = lambda x: [y for y in x] # 按字符分割
# 加载或创建词汇表
if os.path.exists(config.vocab_path):
vocab = pkl.load(open(config.vocab_path, 'rb')) # 从文件加载
else:
vocab = build_vocab(config.train_path, tokenizer, MAX_VOCAB_SIZE, min_freq=1)
pkl.dump(vocab, open(config.vocab_path, 'wb')) # 保存词汇表
print(f"Vocab size: {len(vocab)}")
# -----------------------n-gram哈希函数定义-----------------------
def biGramHash(sequence, t, buckets):
"""计算第t个位置的bigram哈希值
公式: (前一个词的哈希值 * 质数) % 桶大小
说明:如果t-1越界,用0代替(相当于用PAD的哈希值)
"""
t1 = sequence[t - 1] if t - 1 >= 0 else 0
return (t1 * 14918087) % buckets # 14918087是一个大质数
def triGramHash(sequence, t, buckets):
"""计算第t个位置的trigram哈希值
公式: (前两个词的哈希值组合 * 质数) % 桶大小
说明:如果t-1或t-2越界,用0代替
"""
t1 = sequence[t - 1] if t - 1 >= 0 else 0
t2 = sequence[t - 2] if t - 2 >= 0 else 0
return (t2 * 14918087 * 18408749 + t1 * 14918087) % buckets # 双质数减少冲突
# -------------------------------------------------------------
def load_dataset(path, pad_size=32):
"""加载并处理单个数据集文件
Args:
path: 数据集文件路径
pad_size: 填充/截断后的固定长度
Returns:
contents: 处理后的数据列表,元素为(词索引, 标签, 长度, bigram, trigram)
"""
contents = []
with open(path, 'r', encoding='UTF-8') as f:
for line in tqdm(f): # 显示进度条
lin = line.strip()
if not lin:
continue
content, label = lin.split('\t') # 分割文本和标签
# 分词并处理长度
token = tokenizer(content)
seq_len = len(token)
# 填充或截断至固定长度pad_size
if pad_size:
if len(token) < pad_size:
# 注意:这里用vocab.get(PAD)可能存在错误,PAD应为词汇表中已存在的键
token.extend([vocab.get(PAD)] * (pad_size - len(token))) # 填充
else:
token = token[:pad_size] # 截断
seq_len = pad_size # 更新实际长度为pad_size
# 将词转换为索引,未知词用UNK的索引
words_line = []
for word in token:
words_line.append(vocab.get(word, vocab.get(UNK))) # 双重保险取UNK索引
# 生成n-gram特征(FastText模型需要)
buckets = config.n_gram_vocab # 从配置获取哈希桶数量
bigram = []
trigram = []
for i in range(pad_size):
# 为每个位置生成bigram和trigram的哈希值
bigram.append(biGramHash(words_line, i, buckets))
trigram.append(triGramHash(words_line, i, buckets))
# 添加处理后的数据:词索引、标签、长度、bigram、trigram
contents.append((words_line, int(label), seq_len, bigram, trigram))
return contents # 返回结构:[([...], 0, 32, [...], [...]), ...]
# 加载并处理所有数据集
train = load_dataset(config.train_path, config.pad_size)
dev = load_dataset(config.dev_path, config.pad_size)
test = load_dataset(config.test_path, config.pad_size)
return vocab, train, dev, test
字/词分割
python
if use_word:
tokenizer = lambda x: x.split(' ') # 按空格分割(词级别)
else:
tokenizer = lambda x: [y for y in x] # 按字符分割
首先定义分词器,如果是按照单词分割,就按照空格做分割;如果是按照字分割,就按照字符做分割。代码里的lambda表达式可以换成常规的写法:
python
if use_word:
def tokenizer(x):
return x.split(' ') # 按空格分割成词列表
else:
def tokenizer(x):
return [y for y in x] # 拆分成字符列表
加载/创建词汇表
接下来我们需要得到一个词汇表,它的作用就是把文字转换成数字ID。这里可以加载现成的,也可以自己生成一个词汇表。
python
# 加载或创建词汇表
if os.path.exists(config.vocab_path):
vocab = pkl.load(open(config.vocab_path, 'rb')) # 从文件加载
else:
vocab = build_vocab(config.train_path, tokenizer, MAX_VOCAB_SIZE, min_freq=1)
pkl.dump(vocab, open(config.vocab_path, 'wb')) # 保存词汇表
print(f"Vocab size: {len(vocab)}")
构建词汇表的第一步是计算各个文字出现的频次。
python
vocab_dic[word] = vocab_dic.get(word, 0) + 1
我们遍历每一行文本内容,将每一句文本的回车符去掉,并按照之前设计好的分词器进行分割。

遍历之后可以得到一个字典,字典里面记录的是每个字出现的次数。

这行代码的作用是从词汇字典中筛选出符合最小词频要求的单词,并按词频从高到低排序,最后截取前 max_size个单词形成最终的词汇列表。
python
vocab_list = sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq], key=lambda x: x[1], reverse=True)[:max_size]
-
vocab_dic.items()
获取字典中的键值对列表,格式为
[(单词1, 词频1), (单词2, 词频2), ...]
。示例输入:
{'apple':5, 'banana':3, 'cherry':7} → [('apple',5), ('banana',3), ('cherry',7)]
-
列表推导式筛选
if _ >= min_freq
过滤出词频≥min_freq 的单词,
_
表示元组的第二个元素(词频)。如果min_freq=4, 筛选后:[('apple',5), ('cherry',7)]
(banana
因词频3被剔除) -
**按词频降序排序
sorted(..., key=lambda x: x, reverse=True)
**key=lambda x: x
:指定按元组的第二个元素(词频)排序。reverse=True
:降序排列(从高到低)。
排序后:[('cherry',7), ('apple',5)]
-
截取前
max_size
个元素[:max_size]
保留排序后的前
max_size
个高频词。若
max_size=1
:→ 结果:
[('cherry',7)]
-
最终输出
vocab_list
得到处理后的词汇列表,格式为
[(单词, 词频), ...]
,按词频降序排列且长度≤max_size
。

得到上图所示的文字频次表后,再做一步处理。下面这行代码的核心作用 是将排序后的词汇列表转换为 **{字: 索引}
**的字典映射。
python
vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}
这行代码我看着也头大,可以拆解成下面这种写法:
python
vocab_dic = {}
for idx, word_count in enumerate(vocab_list):
word = word_count # word_count 是 (单词, 词频) 元组,取第一个元素即单词
vocab_dic[word] = idx
上面的代码看起来就清晰多了,先是将原先的数组转换为元组,这样每个字都被赋予了一个数字ID,再将{字:数字ID}的形式存到哈希表里,如下图所示:

python
# 添加未知词和填充符的索引(排在最后两位)
vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})
最后再把未知词和填充符号添加到字典的末尾就好了。
数据集处理
python
with open(path, 'r', encoding='UTF-8') as f:
for line in tqdm(f): # 显示进度条
lin = line.strip()
if not lin:
continue
content, label = lin.split('\t') # 分割文本和标签
# 分词并处理长度
token = tokenizer(content)
seq_len = len(token)
# 填充或截断至固定长度pad_size
if pad_size:
if len(token) < pad_size:
# 注意:这里用vocab.get(PAD)可能存在错误,PAD应为词汇表中已存在的键
token.extend([vocab.get(PAD)] * (pad_size - len(token))) # 填充
else:
token = token[:pad_size] # 截断
seq_len = pad_size # 更新实际长度为pad_size
# 将词转换为索引,未知词用UNK的索引
words_line = []
for word in token:
words_line.append(vocab.get(word, vocab.get(UNK))) # 双重保险取UNK索引
# 生成n-gram特征(FastText模型需要)
buckets = config.n_gram_vocab # 从配置获取哈希桶数量
bigram = []
trigram = []
for i in range(pad_size):
# 为每个位置生成bigram和trigram的哈希值
bigram.append(biGramHash(words_line, i, buckets))
trigram.append(triGramHash(words_line, i, buckets))
# 添加处理后的数据:词索引、标签、长度、bigram、trigram
contents.append((words_line, int(label), seq_len, bigram, trigram))
在训练的时候我们要保证每条语句的长度是一致的,所以要设置一个固定长度pad_size。小于这个长度就在句子后面增加PAD符号,大于这个长度就做后向截断。
python
# 生成n-gram特征(FastText模型需要)
buckets = config.n_gram_vocab # 从配置获取哈希桶数量
bigram = []
trigram = []
for i in range(pad_size):
# 为每个位置生成bigram和trigram的哈希值
bigram.append(biGramHash(words_line, i, buckets))
trigram.append(triGramHash(words_line, i, buckets))
这段代码通过生成 Bigram(二元组) 和 Trigram(三元组) 的哈希特征,为深度学习模型提供局部词序信息,增强模型对短语和上下文关系的捕捉能力,尤其在处理短文本时效果显著。
- Bigram:相邻两个词的组合(如"深度学习" → "深度"-"学习")
- Trigram:相邻三个词的组合(如"自然语言处理" → "自然"-"语言"-"处理")
在实际的实现里我们做了哈希映射,原因是直接存储所有可能的N-Gram会导致特征维度爆炸。哈希映射的原理是将任意长度的N-Gram映射到固定范围的桶,做维度压缩。
python
def biGramHash(sequence, t, buckets):
"""计算第t个位置的bigram哈希值
公式: (前一个词的哈希值 * 质数) % 桶大小
说明:如果t-1越界,用0代替(相当于用PAD的哈希值)
"""
t1 = sequence[t - 1] if t - 1 >= 0 else 0
return (t1 * 14918087) % buckets # 14918087是一个大质数
def triGramHash(sequence, t, buckets):
"""计算第t个位置的trigram哈希值
公式: (前两个词的哈希值组合 * 质数) % 桶大小
说明:如果t-1或t-2越界,用0代替
"""
t1 = sequence[t - 1] if t - 1 >= 0 else 0
t2 = sequence[t - 2] if t - 2 >= 0 else 0
return (t2 * 14918087 * 18408749 + t1 * 14918087) % buckets # 双质数减少冲突
使用大质数组合相乘的设计,主要目的是通过数学特性降低哈希冲突率,同时保证计算效率。
质数的特点是只能被1和自身整除,在乘法运算中不同质数组合能生成唯一性更高的中间值;另外,使用两个间距大的千万级大质数,能够避免相邻词索引的小幅变化导致哈希值相似。比如:
python
词t-2=100, 词t-1=101 → 100*18,408,749 + 101*14,918,087 ≈ 3.3e9
词t-2=101, 词t-1=100 → 101*18,408,749 + 100*14,918,087 ≈ 3.4e9