sundea/text-classification

1# coding: UTF-82import os3import torch4import numpy as np5import pickle as pkl6from tqdm import tqdm7import time8from datetime import timedelta9 10 11MAX_VOCAB_SIZE = 10000  # 词表长度限制12UNK, PAD = '<UNK>', '<PAD>'  # 未知字，padding符号13 14 15def build_vocab(file_path, tokenizer, max_size, min_freq):16    vocab_dic = {}17    with open(file_path, 'r', encoding='UTF-8') as f:18        for line in tqdm(f):19            lin = line.strip()20            if not lin:21                continue22            content = lin.split('\t')[0]23            for word in tokenizer(content):24                vocab_dic[word] = vocab_dic.get(word, 0) + 125        vocab_list = sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq], key=lambda x: x[1], reverse=True)[:max_size]26        vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}27        vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})28    return vocab_dic29 30 31def build_dataset(config, ues_word):32    if ues_word:33        tokenizer = lambda x: x.split(' ')  # 以空格隔开，word-level34    else:35        tokenizer = lambda x: [y for y in x]  # char-level36    if os.path.exists(config.vocab_path):37        vocab = pkl.load(open(config.vocab_path, 'rb'))38    else:39        vocab = build_vocab(config.train_path, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)40        pkl.dump(vocab, open(config.vocab_path, 'wb'))41    print(f"Vocab size: {len(vocab)}")42 43    def load_dataset(path, pad_size=32):44        contents = []45        with open(path, 'r', encoding='UTF-8') as f:46            for line in tqdm(f):47                lin = line.strip()48                if not lin:49                    continue50                content, label = lin.split('\t')51                words_line = []52                token = tokenizer(content)53                seq_len = len(token)54                if pad_size:55                    if len(token) < pad_size:56                        token.extend([PAD] * (pad_size - len(token)))57                    else:58                        token = token[:pad_size]59                        seq_len = pad_size60                # word to id61                for word in token:62                    words_line.append(vocab.get(word, vocab.get(UNK)))63                contents.append((words_line, int(label), seq_len))64        return contents  # [([...], 0), ([...], 1), ...]65    train = load_dataset(config.train_path, config.pad_size)66    dev = load_dataset(config.dev_path, config.pad_size)67    test = load_dataset(config.test_path, config.pad_size)68    return vocab, train, dev, test69 70 71class DatasetIterater(object):72    def __init__(self, batches, batch_size, device):73        self.batch_size = batch_size74        self.batches = batches75        self.n_batches = len(batches) // batch_size76        self.residue = False  # 记录batch数量是否为整数77        if len(batches) % self.n_batches != 0:78            self.residue = True79        self.index = 080        self.device = device81 82    def _to_tensor(self, datas):83        x = torch.LongTensor([_[0] for _ in datas]).to(self.device)84        y = torch.LongTensor([_[1] for _ in datas]).to(self.device)85 86        # pad前的长度(超过pad_size的设为pad_size)87        seq_len = torch.LongTensor([_[2] for _ in datas]).to(self.device)88        return (x, seq_len), y89 90    def __next__(self):91        if self.residue and self.index == self.n_batches:92            batches = self.batches[self.index * self.batch_size: len(self.batches)]93            self.index += 194 95            batches = self._to_tensor(batches)96            return batches97 98        elif self.index >= self.n_batches:99            self.index = 0100            raise StopIteration101        else:102            batches = self.batches[self.index * self.batch_size: (self.index + 1) * self.batch_size]103            self.index += 1104            batches = self._to_tensor(batches)105            return batches106 107    def __iter__(self):108        return self109 110    def __len__(self):111        if self.residue:112            return self.n_batches + 1113        else:114            return self.n_batches115 116 117def build_iterator(dataset, config):118    iter = DatasetIterater(dataset, config.batch_size, config.device)119    return iter120 121 122def get_time_dif(start_time):123    """获取已使用时间"""124    end_time = time.time()125    time_dif = end_time - start_time126    return timedelta(seconds=int(round(time_dif)))127 128 129if __name__ == "__main__":130    '''提取预训练词向量'''131    # 下面的目录、文件名按需更改。132    train_dir = "./THUCNews/data/train.txt"133    vocab_dir = "./THUCNews/data/vocab.pkl"134    pretrain_dir = "./THUCNews/data/sgns.sogou.char"135    emb_dim = 300136    filename_trimmed_dir = "./THUCNews/data/embedding_SougouNews"137    if os.path.exists(vocab_dir):138        word_to_id = pkl.load(open(vocab_dir, 'rb'))139    else:140        # tokenizer = lambda x: x.split(' ')  # 以词为单位构建词表(数据集中词之间以空格隔开)141        tokenizer = lambda x: [y for y in x]  # 以字为单位构建词表142        word_to_id = build_vocab(train_dir, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)143        pkl.dump(word_to_id, open(vocab_dir, 'wb'))144 145    embeddings = np.random.rand(len(word_to_id), emb_dim)146    f = open(pretrain_dir, "r", encoding='UTF-8')147    for i, line in enumerate(f.readlines()):148        # if i == 0:  # 若第一行是标题，则跳过149        #     continue150        lin = line.strip().split(" ")151        if lin[0] in word_to_id:152            idx = word_to_id[lin[0]]153            emb = [float(x) for x in lin[1:301]]154            embeddings[idx] = np.asarray(emb, dtype='float32')155    f.close()156    np.savez_compressed(filename_trimmed_dir, embeddings=embeddings)157