load_data.py 3.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  1. import os
  2. import re
  3. import jieba
  4. import time
  5. import numpy as np
  6. # jieba.enable_parallel() # jieba支持多进程 ,但是win不支持
  7. token = "[0-9\s+\.\!\/_,$%^*()?;;:【】+\"\'\[\]\\]+|[+——!,;:。?《》、~@#¥%……&*()“”.=-]+"
  8. labels_index = {} # 记录分类标签的序号
  9. stopwords = set(open('dict/stop_words.txt', encoding='utf-8').read().split()) # 停用词
  10. # for scikit part
  11. def preprocess(text):
  12. text1 = re.sub('&nbsp', ' ', text)
  13. str_no_punctuation = re.sub(token, ' ', text1) # 去掉标点
  14. text_list = list(jieba.cut(str_no_punctuation)) # 分词列表
  15. text_list = [item for item in text_list if item != ' '] # 去掉空格
  16. return ' '.join(text_list)
  17. def load_datasets():
  18. # should run corpus_split.py first
  19. base_dir = 'tag/'
  20. X_data = {'train':[], 'test':[]}
  21. y = {'train':[], 'test':[]}
  22. for type_name in ['train', 'test']:
  23. corpus_dir = os.path.join(base_dir, type_name)
  24. for label in os.listdir(corpus_dir):
  25. label_dir = os.path.join(corpus_dir, label)
  26. file_list = os.listdir(label_dir)
  27. print("label: {}, len: {}".format(label, len(file_list)))
  28. for fname in file_list:
  29. file_path = os.path.join(label_dir, fname)
  30. with open(file_path, encoding='UTF-8', errors='ignore') as text_file:
  31. text_content = preprocess(text_file.read())
  32. X_data[type_name].append(text_content)
  33. y[type_name].append(label)
  34. print("{} corpus len: {}\n".format(type_name, len(X_data[type_name])))
  35. return X_data['train'], y['train'], X_data['test'], y['test']
  36. # for keras part
  37. def preprocess_keras(text):
  38. text1 = re.sub('&nbsp', ' ', text)
  39. str_no_punctuation = re.sub(token, ' ', text1) # 去掉标点
  40. text_list = list(jieba.cut(str_no_punctuation)) # 分词列表
  41. text_list = [item for item in text_list if item != ' ' and item not in stopwords] # 去掉空格和停用词
  42. return ' '.join(text_list)
  43. def load_raw_datasets():
  44. labels = []
  45. texts = []
  46. base_dir = 'CN_Corpus/SogouC.reduced/Reduced'
  47. t1 = time.time()
  48. for cate_index, label in enumerate(os.listdir(base_dir)):
  49. label_dir = os.path.join(base_dir, label)
  50. file_list = os.listdir(label_dir)
  51. labels_index[label] = cate_index # 记录分类标签的整数标号
  52. print("label: {}, len: {}".format(label, len(file_list)))
  53. for fname in file_list:
  54. f = open(os.path.join(label_dir, fname), encoding='gb2312', errors='ignore')
  55. texts.append(preprocess_keras(f.read()))
  56. f.close()
  57. labels.append(labels_index[label])
  58. t2 = time.time()
  59. tm_cost = t2-t1
  60. print('\nDone. {} total categories, {} total docs. cost {} seconds.'.format(len(os.listdir(base_dir)), len(texts), tm_cost))
  61. return texts, labels
  62. def load_pre_trained():
  63. # load pre-trained embedding model
  64. embeddings_index = {}
  65. with open('Embedding/sgns.sogou.word') as f:
  66. _, embedding_dim = f.readline().split()
  67. for line in f:
  68. values = line.split()
  69. word = values[0]
  70. coefs = np.asarray(values[1:], dtype='float32')
  71. embeddings_index[word] = coefs
  72. print('Found %s word vectors, dimension %s' % (len(embeddings_index), embedding_dim))
  73. return embeddings_index