hoodiexxx/Bert_Chinese_Text_Classification_Model

1from torch.utils.data import Dataset2import random3import numpy as np4from tqdm import tqdm5import torch6 7import sen2inds8 9 10class textCNN_data(Dataset):11    def __init__(self, trainDataFile):12        trainData = open(trainDataFile, 'r').read().split('\n')13        trainData = list(filter(None, trainData))14 15        res = []16        for data in tqdm(trainData, desc='index to tensor'):17            data = list(filter(None, data.split(',')))18            data = [int(x) for x in data]19            cla = torch.tensor(data[0], dtype=torch.long)20            sentence = torch.tensor(data[1:], dtype=torch.long)21            temp = []22            temp.append(cla)23            temp.append(sentence)24            res.append(temp)25 26        self.trainData = res27 28    def __len__(self):29        return len(self.trainData)30 31    def __getitem__(self, idx):32        data = self.trainData[idx]33        cla = data[0]34        sentence = data[1]35 36        return cla, sentence37 38 39word2ind, ind2word = sen2inds.get_worddict('wordLabel.txt')40label_w2n, label_n2w = sen2inds.read_labelFile('data/label2.txt') # sen2inds.read_labelFile('data/label2.txt')41 42textCNN_param = {43    'vocab_size': len(word2ind) + 1, # plus one for 0 padding44    'embed_dim': 256,  # 1 x 128 vector45    'class_num': len(label_w2n),46    "kernel_num": 16,47    "kernel_size": [3, 4, 5],48    "dropout": 0.5,49}50dataLoader_param = {51    'batch_size': 128,52    'shuffle': True,53}54