| 12345678910111213141516171819202122232425262728293031323334 |
- from torchtext.vocab import build_vocab_from_iterator
- import torch
- from torch.nn.utils.rnn import pad_sequence
- #实现build_vocab函数,基于dataset,建立词汇表
- def build_vocab(dataset):
- #unk表示未知词,pad表示填充词
- special=['<unk>','<pad>']
- text_iter=map(lambda x:x[0] , dataset)#文本序列
- #建立文本词汇表text_vocab
- #将min_freq设置为2,也就是至少出现两次的单词,才会添加到词表中
- text_vocab = build_vocab_from_iterator(text_iter,min_freq=2,specials=special)
- #将unk对应的索引,设置为默认索引
- text_vocab.set_default_index(text_vocab['<unk>'])
- return text_vocab
- #整个collate_batch函数是批量读取文本数据时的回调函数
- def collate_batch(batch,text_vocab):
- text_list=list()
- labels=list()
- #每次读取一组数据
- for text,label in batch:
- text_tokens = [text_vocab[token] for token in text]
- text_tensor=torch.tensor(text_tokens,dtype=torch.long)
- text_list.append(text_tensor)
- labels.append(torch.tensor(label - 1,dtype=torch.long))
-
- padding_idx = text_vocab["<pad>"]
- #将batch填充为相同长度文本
- text_padded = pad_sequence(text_list,batch_first=True,padding_value=padding_idx)
- #一次性堆叠成批次 形状: (N, 3)
- labels_tensor=torch.stack(labels)
- #返回文本和标签的张量形式，用于后续的模型训练
- return text_padded,labels_tensor
|
