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