-
Notifications
You must be signed in to change notification settings - Fork 115
/
bilsm_crf_model.py
27 lines (24 loc) · 924 Bytes
/
bilsm_crf_model.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from keras.models import Sequential
from keras.layers import Embedding, Bidirectional, LSTM
from keras_contrib.layers import CRF
import process_data
import pickle
EMBED_DIM = 200
BiRNN_UNITS = 200
def create_model(train=True):
if train:
(train_x, train_y), (test_x, test_y), (vocab, chunk_tags) = process_data.load_data()
else:
with open('model/config.pkl', 'rb') as inp:
(vocab, chunk_tags) = pickle.load(inp)
model = Sequential()
model.add(Embedding(len(vocab), EMBED_DIM, mask_zero=True)) # Random embedding
model.add(Bidirectional(LSTM(BiRNN_UNITS // 2, return_sequences=True)))
crf = CRF(len(chunk_tags), sparse_target=True)
model.add(crf)
model.summary()
model.compile('adam', loss=crf.loss_function, metrics=[crf.accuracy])
if train:
return model, (train_x, train_y), (test_x, test_y)
else:
return model, (vocab, chunk_tags)