|
| 1 | +import pickle |
| 2 | +import numpy as np |
| 3 | +import platform |
| 4 | +from collections import Counter |
| 5 | + |
| 6 | +from keras.models import Sequential |
| 7 | +from keras.layers import Embedding, Bidirectional, LSTM |
| 8 | +from keras_contrib.layers import CRF |
| 9 | +from keras.preprocessing.sequence import pad_sequences |
| 10 | + |
| 11 | +EMBED_DIM = 200 |
| 12 | +BiRNN_UNITS = 200 |
| 13 | + |
| 14 | + |
| 15 | + |
| 16 | +def load_data(): |
| 17 | + train = _parse_data(open('zh-NER/data/train_data.data', 'rb')) |
| 18 | + test = _parse_data(open('zh-NER/data/test_data.data', 'rb')) |
| 19 | + |
| 20 | + word_counts = Counter(row[0].lower() for sample in train for row in sample) |
| 21 | + vocab = [w for w, f in iter(word_counts.items()) if f >= 2] |
| 22 | + chunk_tags = ['O', 'B-PER', 'I-PER', 'B-LOC', 'I-LOC', "B-ORG", "I-ORG"] |
| 23 | + |
| 24 | + # save initial config data |
| 25 | + with open('zh-NER/model/config.pkl', 'wb') as outp: |
| 26 | + pickle.dump((vocab, chunk_tags), outp) |
| 27 | + |
| 28 | + train = _process_data(train, vocab, chunk_tags) |
| 29 | + test = _process_data(test, vocab, chunk_tags) |
| 30 | + return train, test, (vocab, chunk_tags) |
| 31 | + |
| 32 | + |
| 33 | +def _parse_data(fh): |
| 34 | + # in windows the new line is '\r\n\r\n' the space is '\r\n' . so if you use windows system, |
| 35 | + # you have to use recorsponding instructions |
| 36 | + |
| 37 | + if platform.system() == 'Windows': |
| 38 | + split_text = '\r\n' |
| 39 | + else: |
| 40 | + split_text = '\n' |
| 41 | + |
| 42 | + string = fh.read().decode('utf-8') |
| 43 | + data = [[row.split() for row in sample.split(split_text)] for |
| 44 | + sample in |
| 45 | + string.strip().split(split_text + split_text)] |
| 46 | + fh.close() |
| 47 | + return data |
| 48 | + |
| 49 | + |
| 50 | +def _process_data(data, vocab, chunk_tags, maxlen=None, onehot=False): |
| 51 | + if maxlen is None: |
| 52 | + maxlen = max(len(s) for s in data) |
| 53 | + word2idx = dict((w, i) for i, w in enumerate(vocab)) |
| 54 | + x = [[word2idx.get(w[0].lower(), 1) for w in s] for s in data] # set to <unk> (index 1) if not in vocab |
| 55 | + |
| 56 | + y_chunk = [[chunk_tags.index(w[1]) for w in s] for s in data] |
| 57 | + |
| 58 | + x = pad_sequences(x, maxlen) # left padding |
| 59 | + |
| 60 | + y_chunk = pad_sequences(y_chunk, maxlen, value=-1) |
| 61 | + |
| 62 | + if onehot: |
| 63 | + y_chunk = np.eye(len(chunk_tags), dtype='float32')[y_chunk] |
| 64 | + else: |
| 65 | + y_chunk = np.expand_dims(y_chunk, 2) |
| 66 | + return x, y_chunk |
| 67 | + |
| 68 | + |
| 69 | +def process_data(data, vocab, maxlen=100): |
| 70 | + word2idx = dict((w, i) for i, w in enumerate(vocab)) |
| 71 | + x = [word2idx.get(w[0].lower(), 1) for w in data] |
| 72 | + length = len(x) |
| 73 | + x = pad_sequences([x], maxlen) # left padding |
| 74 | + return x, length |
| 75 | + |
| 76 | + |
| 77 | +def create_model(train=True): |
| 78 | + if train: |
| 79 | + (train_x, train_y), (test_x, test_y), (vocab, chunk_tags) = load_data() |
| 80 | + else: |
| 81 | + with open('model/config.pkl', 'rb') as inp: |
| 82 | + (vocab, chunk_tags) = pickle.load(inp) |
| 83 | + model = Sequential() |
| 84 | + model.add(Embedding(len(vocab), EMBED_DIM, mask_zero=True)) # Random embedding |
| 85 | + model.add(Bidirectional(LSTM(BiRNN_UNITS // 2, return_sequences=True))) |
| 86 | + crf = CRF(len(chunk_tags), sparse_target=True) |
| 87 | + model.add(crf) |
| 88 | + model.summary() |
| 89 | + model.compile('adam', loss=crf.loss_function, metrics=[crf.accuracy]) |
| 90 | + if train: |
| 91 | + return model, (train_x, train_y), (test_x, test_y) |
| 92 | + else: |
| 93 | + return model, (vocab, chunk_tags) |
| 94 | + |
| 95 | + |
| 96 | +def train(): |
| 97 | + EPOCHS = 10 |
| 98 | + model, (train_x, train_y), (test_x, test_y) = create_model() |
| 99 | + # train model |
| 100 | + model.fit(train_x, train_y,batch_size=16,epochs=EPOCHS, validation_data=[test_x, test_y]) |
| 101 | + model.save('model/crf.h5') |
| 102 | + |
| 103 | +def test(): |
| 104 | + model, (vocab, chunk_tags) = create_model(train=False) |
| 105 | + predict_text = '中华人民共和国国务院总理周恩来在外交部长陈毅的陪同下，连续访问了埃塞俄比亚等非洲10国以及阿尔巴尼亚' |
| 106 | + str, length = process_data(predict_text, vocab) |
| 107 | + model.load_weights('model/crf.h5') |
| 108 | + raw = model.predict(str)[0][-length:] |
| 109 | + result = [np.argmax(row) for row in raw] |
| 110 | + result_tags = [chunk_tags[i] for i in result] |
| 111 | + |
| 112 | + per, loc, org = '', '', '' |
| 113 | + |
| 114 | + for s, t in zip(predict_text, result_tags): |
| 115 | + if t in ('B-PER', 'I-PER'): |
| 116 | + per += ' ' + s if (t == 'B-PER') else s |
| 117 | + if t in ('B-ORG', 'I-ORG'): |
| 118 | + org += ' ' + s if (t == 'B-ORG') else s |
| 119 | + if t in ('B-LOC', 'I-LOC'): |
| 120 | + loc += ' ' + s if (t == 'B-LOC') else s |
| 121 | + |
| 122 | + print(['person:' + per, 'location:' + loc, 'organzation:' + org]) |
| 123 | + |
| 124 | + |
| 125 | +if __name__ == "__main__": |
| 126 | + train() |
0 commit comments