清華大學THUCNews數據集處理方法

清華大學THUCNews數據集處理方法

數據集下載地址:http://thuctc.thunlp.org/

本文主要記錄了清華的THUCNEWS數據集用於文本分類的數據預處理方法，比較簡單直接上代碼

#! usr/bin/env python3
# -*- coding:utf-8 -*-

"""
 清華大學的文本分類數據集的處理
 @Author：MaCan
 @Time：2019/9/17 11:40
 @File：thunews_process.py
 @Mail：[email protected]
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import re
import pickle


def load_file(path):
    class_file_num = {}
    with open('thu_news_3w_v1.txt', 'w', encoding='utf-8') as fw:
        for sub_dir in os.listdir(path):
            print(sub_dir)
            if sub_dir in ['星座']:
                continue
            curr_dir = os.path.join(path, sub_dir)
            if os.path.isdir(curr_dir):
                for file in os.listdir(curr_dir):
                    curr_file = os.path.join(curr_dir, file)
                    # print(curr_file)
                    context = ''
                    with open(curr_file, 'r', encoding='utf-8') as fd:
                        idx = 0
                        for line in fd:
                            if idx == 0 and len(line) < 30:
                                context += line + '|||'
                            else:
                                context += line
                            idx += 1
                    if len(context) < 10:
                        continue
                    context = re.sub('(\ud83c[\udf00-\udfff])|(\ud83d[\udc00-\ude4f\ude80-\udeff])|[\u2600-\u2B55]', '',
                                     context.strip())
                    context = re.sub('[\n\r\t  　　\\\]+', '', context)
                    context = '__' + sub_dir + '__\t' + context
                    # print(context)
                    fw.write(context + '\n')
                    class_file_num[sub_dir] = class_file_num.get(sub_dir, 0) + 1
                    if class_file_num.get(sub_dir) > 30000:
                        break
    print(class_file_num)


def cut_sentence(line, max_seq_len, mode):
    """
    文本截斷
    :param line:
    :param mode 截斷方式：[扔掉末尾，扔掉開始，title+top+tail, top+tail-title]
    :return:
    """
    line = line.strip()
    if mode == 0:
        if len(line) > max_seq_len:
            line = line[0:max_seq_len]
        return line
    elif mode == 1:
        if len(line) > max_seq_len:
            line = line[:-max_seq_len]
        return line
    else:
        line = line.split('|||')
        context = ''
        if len(line) >= 2:#title+top+tail
            if mode == 2:
                context += line[0]
            # elif mode == 3:
            #     context += line[1:]
            line = ''.join(line[1:])
        else:
            line = ''.join(line)

        # 按照標點符號進行句子切割
        sentences = re.split(r"([.。!！?？；;，,\s+])", line)
        sentences.append("")
        sentences = ["".join(i) for i in zip(sentences[0::2], sentences[1::2])]
        head = []
        tail = []
        if mode == 2:
            head.append(context)

        def check_curr_sen_len():
            if len(''.join(head)) + len(''.join(tail)) > max_seq_len:
                return True
            return False

        for idx in range(len(sentences)):
            if check_curr_sen_len():
                break
            head.append(sentences[idx])
            if check_curr_sen_len():
                break
            tail.insert(0, sentences[-(idx+1)])
        # merge
        return ''.join(head) + ''.join(tail)

# def test_load_file():
#     load_file(path)


if __name__ == '__main__':
    path = 'F:\THUCNews\THUCNews'
    load_file(path)
    mode = 2
    if mode == 0:
        with open('thu_news_3w_v2_cutmode_3.txt', 'w', encoding='utf-8') as fw:
            with open('thu_news_3w_v2.txt', 'r', encoding='utf-8') as fd:
                for line in fd:
                    line = line.split('\t')
                    new_context = cut_sentence(''.join(line[1:]), 200, 3)
                    fw.write(line[0] + '\t' + new_context + '\n')
    elif mode == 1:
        datas = []
        labels = []
        with open('thu_news_3w_v1_cutmode_2.txt', 'r', encoding='utf-8') as fd:
            for line in fd:
                line = line.split('\t')
                if len(line) == 2:
                    labels.append(line[0])
                    datas.append(line[1])
        from sklearn.model_selection import train_test_split
        print('all data size:{}, {}'.format(len(datas), len(labels)))
        train_x, last_x, train_y, last_y = train_test_split(datas, labels, test_size=0.4)
        dev_x, test_x, dev_y, test_y = train_test_split(last_x, last_y, test_size=0.5)

        def save_list(datas ,labels, path):
            print(path + ' data size:{}, {}'.format(len(datas), len(labels)))
            with open(path, 'w', encoding='utf-8') as fd:
                for data, label in zip(datas, labels):
                    fd.write(label + '\t' + data)

        save_list(train_x, train_y, 'train.txt')
        save_list(dev_x, dev_y, 'dev.txt')
        save_list(test_x, test_y, 'test.txt')

    elif mode == 2:
        datas = []
        labels = []
        with open('thu_news_3w_v1_cutmode_2.txt', 'r', encoding='utf-8') as fd:
            for line in fd:
                line = line.split('\t')
                if len(line) == 2:
                    label = re.sub('[_]+', '', line[0])
                    if label in ['體育', '財經'] and len(labels) < 10000:
                        labels.append(line[0])
                        datas.append(line[1])
        from sklearn.model_selection import train_test_split
        print('all data size:{}, {}'.format(len(datas), len(labels)))
        train_x, last_x, train_y, last_y = train_test_split(datas, labels, test_size=0.4)
        dev_x, test_x, dev_y, test_y = train_test_split(last_x, last_y, test_size=0.5)

        def save_list(datas ,labels, path):
            print(path + ' data size:{}, {}'.format(len(datas), len(labels)))
            with open(path, 'w', encoding='utf-8') as fd:
                for data, label in zip(datas, labels):
                    fd.write(label + '\t' + data)

        save_list(train_x, train_y, 'mini_train.txt')
        save_list(dev_x, dev_y, 'mini_dev.txt')
        save_list(test_x, test_y, 'mini_test.txt')

在main中，mode：0 表示將文本進行切分，具體的切分方法在cut_sentence中，mode:1 表示將數據進行訓練集，驗證集合測試集的劃分。