def to_tensor(value, dtype=torch.float32):
if torch.is_tensor(value):
return value.to(conf.device(), dtype=dtype)
if isinstance(value, (list, tuple, np.ndarray)):
return torch.tensor(value, dtype=dtype, device=conf.device())
if isinstance(value, (float, np.float16, np.float32, np.float64)):
return torch.tensor([value], dtype=dtype, device=conf.device())
raise ValueError("Fail to convert {elem} to tensor".format(elem=value))
def init_hidden(self):
bidirect = 2 if self.bidirectional else 1
if self.rnn_type == 'lstm':
return (torch.zeros(self.num_layers * bidirect,
1,
self.hidden_size,
device=conf.device()),
torch.zeros(self.num_layers * bidirect,
1,
self.hidden_size,
device=conf.device()))
elif self.rnn_type == 'gru':
return torch.zeros(self.num_layers * bidirect,
1,
self.hidden_size,
device=conf.device())
def train_and_dump(from_model=None,
optimizer='sgd',
lr=1e-4,
rnn_type='lstm',
lang_pkl='lang.pt',
drop_n=0,
real_corpus_sample=0.3):
if from_model:
model = EntityRecognizer.load(from_model)
model_pkl_name = from_model
else:
model = EntityRecognizer(lang=Lang.load(lang_pkl),
embedding_dim=200,
rnn_type=rnn_type)
model.init_params()
model_pkl_name = 'model.{rnn_type}.{optimizer}'.format(
rnn_type=rnn_type, optimizer=optimizer)
model.move_to_device(conf.device())
dataset = islice(generate_dataset(real_corpus_sample), drop_n, None)
do_train(model, dataset, model_pkl_name, optimizer=optimizer, lr=lr)
def __init__(
self,
lang,
embedding_dim=200,
hidden_size=512,
rnn_type='lstm',
num_layers=2,
):
super(EntityRecognizer, self).__init__()
self.hidden_size = hidden_size
self.input_size = embedding_dim
self.num_layers = num_layers
self.lang = lang
self.embedding = nn.Embedding(num_embeddings=lang.vocab_size(),
embedding_dim=embedding_dim)
self.bidirectional = True
self.rnn_type = rnn_type
assert rnn_type in ('lstm',
'gru'), 'un-support rnn type:{}'.format(rnn_type)
if self.rnn_type == 'lstm':
self.rnn = nn.LSTM(input_size=embedding_dim,
hidden_size=hidden_size,
dropout=0.5,
num_layers=self.num_layers,
bidirectional=self.bidirectional)
elif self.rnn_type == 'gru':
self.rnn = nn.GRU(input_size=embedding_dim,
hidden_size=hidden_size,
dropout=0.5,
num_layers=self.num_layers,
bidirectional=self.bidirectional)
self.hidden2tag = nn.Linear(hidden_size *
2 if self.bidirectional else 1,
out_features=3)
self.hidden = self.init_hidden()
self.move_to_device(conf.device())
def load_model(model_name):
model = EntityRecognizer.load(model_name)
model.eval()
model.move_to_device(conf.device())
return model