wehub-resource-sync
198 lines
8.1 KiB
Python
198 lines
8.1 KiB
Python
# -*- coding:utf-8 -*-
|
|
# Author: hankcs
|
|
# Date: 2019-08-24 21:49
|
|
import os
|
|
from typing import Tuple, Union, List
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
from tensorflow.python.ops import math_ops
|
|
|
|
from hanlp.common.vocab_tf import VocabTF
|
|
from hanlp.utils.io_util import get_resource
|
|
from hanlp.utils.torch_util import load_word2vec
|
|
from hanlp.utils.tf_util import hanlp_register
|
|
from hanlp_common.util import DummyContext
|
|
|
|
|
|
class Word2VecEmbeddingV1(tf.keras.layers.Layer):
|
|
def __init__(self, path: str = None, vocab: VocabTF = None, normalize: bool = False, load_all=True, mask_zero=True,
|
|
trainable=False, name=None, dtype=None, dynamic=False, **kwargs):
|
|
super().__init__(trainable, name, dtype, dynamic, **kwargs)
|
|
if load_all and vocab and vocab.locked:
|
|
vocab.unlock()
|
|
self.vocab, self.array_np = self._load(path, vocab, normalize)
|
|
self.vocab.lock()
|
|
self.array_ks = tf.keras.layers.Embedding(input_dim=len(self.vocab), output_dim=self.dim, trainable=trainable,
|
|
embeddings_initializer=tf.keras.initializers.Constant(self.array_np),
|
|
mask_zero=mask_zero)
|
|
self.mask_zero = mask_zero
|
|
self.supports_masking = mask_zero
|
|
|
|
def compute_mask(self, inputs, mask=None):
|
|
if not self.mask_zero:
|
|
return None
|
|
|
|
return math_ops.not_equal(inputs, self.vocab.pad_idx)
|
|
|
|
def call(self, inputs, **kwargs):
|
|
return self.array_ks(inputs, **kwargs)
|
|
|
|
def compute_output_shape(self, input_shape):
|
|
return input_shape[0], self.dim
|
|
|
|
@staticmethod
|
|
def _load(path, vocab, normalize=False) -> Tuple[VocabTF, Union[np.ndarray, None]]:
|
|
if not vocab:
|
|
vocab = VocabTF()
|
|
if not path:
|
|
return vocab, None
|
|
assert vocab.unk_idx is not None
|
|
|
|
word2vec, dim = load_word2vec(path)
|
|
for word in word2vec:
|
|
vocab.get_idx(word)
|
|
|
|
pret_embs = np.zeros(shape=(len(vocab), dim), dtype=np.float32)
|
|
state = np.random.get_state()
|
|
np.random.seed(0)
|
|
bias = np.random.uniform(low=-0.001, high=0.001, size=dim).astype(dtype=np.float32)
|
|
scale = np.sqrt(3.0 / dim)
|
|
for word, idx in vocab.token_to_idx.items():
|
|
vec = word2vec.get(word, None)
|
|
if vec is None:
|
|
vec = word2vec.get(word.lower(), None)
|
|
# if vec is not None:
|
|
# vec += bias
|
|
if vec is None:
|
|
# vec = np.random.uniform(-scale, scale, [dim])
|
|
vec = np.zeros([dim], dtype=np.float32)
|
|
pret_embs[idx] = vec
|
|
# noinspection PyTypeChecker
|
|
np.random.set_state(state)
|
|
return vocab, pret_embs
|
|
|
|
@property
|
|
def size(self):
|
|
if self.array_np is not None:
|
|
return self.array_np.shape[0]
|
|
|
|
@property
|
|
def dim(self):
|
|
if self.array_np is not None:
|
|
return self.array_np.shape[1]
|
|
|
|
@property
|
|
def shape(self):
|
|
if self.array_np is None:
|
|
return None
|
|
return self.array_np.shape
|
|
|
|
def get_vector(self, word: str) -> np.ndarray:
|
|
assert self.array_np is not None
|
|
return self.array_np[self.vocab.get_idx_without_add(word)]
|
|
|
|
def __getitem__(self, word: Union[str, List, tf.Tensor]) -> np.ndarray:
|
|
if isinstance(word, str):
|
|
return self.get_vector(word)
|
|
elif isinstance(word, list):
|
|
vectors = np.zeros(shape=(len(word), self.dim))
|
|
for idx, token in enumerate(word):
|
|
vectors[idx] = self.get_vector(token)
|
|
return vectors
|
|
elif isinstance(word, tf.Tensor):
|
|
if word.dtype == tf.string:
|
|
word_ids = self.vocab.token_to_idx_table.lookup(word)
|
|
return tf.nn.embedding_lookup(self.array_tf, word_ids)
|
|
elif word.dtype == tf.int32 or word.dtype == tf.int64:
|
|
return tf.nn.embedding_lookup(self.array_tf, word)
|
|
|
|
|
|
@hanlp_register
|
|
class Word2VecEmbeddingTF(tf.keras.layers.Embedding):
|
|
|
|
def __init__(self, filepath: str = None, vocab: VocabTF = None, expand_vocab=True, lowercase=True,
|
|
input_dim=None, output_dim=None, unk=None, normalize=False,
|
|
embeddings_initializer='VarianceScaling',
|
|
embeddings_regularizer=None,
|
|
activity_regularizer=None, embeddings_constraint=None, mask_zero=True, input_length=None,
|
|
name=None, cpu=True, **kwargs):
|
|
filepath = get_resource(filepath)
|
|
word2vec, _output_dim = load_word2vec(filepath)
|
|
if output_dim:
|
|
assert output_dim == _output_dim, f'output_dim = {output_dim} does not match {filepath}'
|
|
output_dim = _output_dim
|
|
# if the `unk` token exists in the pretrained,
|
|
# then replace it with a self-defined one, usually the one in word vocab
|
|
if unk and unk in word2vec:
|
|
word2vec[vocab.safe_unk_token] = word2vec.pop(unk)
|
|
if vocab is None:
|
|
vocab = VocabTF()
|
|
vocab.update(word2vec.keys())
|
|
if expand_vocab and vocab.mutable:
|
|
for word in word2vec:
|
|
vocab.get_idx(word.lower() if lowercase else word)
|
|
if input_dim:
|
|
assert input_dim == len(vocab), f'input_dim = {input_dim} does not match {filepath}'
|
|
input_dim = len(vocab)
|
|
# init matrix
|
|
self._embeddings_initializer = embeddings_initializer
|
|
embeddings_initializer = tf.keras.initializers.get(embeddings_initializer)
|
|
with tf.device('cpu:0') if cpu else DummyContext():
|
|
pret_embs = embeddings_initializer(shape=[input_dim, output_dim]).numpy()
|
|
# insert to pret_embs
|
|
for word, idx in vocab.token_to_idx.items():
|
|
vec = word2vec.get(word, None)
|
|
# Retry lower case
|
|
if vec is None and lowercase:
|
|
vec = word2vec.get(word.lower(), None)
|
|
if vec is not None:
|
|
pret_embs[idx] = vec
|
|
if normalize:
|
|
pret_embs /= np.std(pret_embs)
|
|
if not name:
|
|
name = os.path.splitext(os.path.basename(filepath))[0]
|
|
super().__init__(input_dim, output_dim, tf.keras.initializers.Constant(pret_embs), embeddings_regularizer,
|
|
activity_regularizer, embeddings_constraint, mask_zero, input_length, name=name, **kwargs)
|
|
self.filepath = filepath
|
|
self.expand_vocab = expand_vocab
|
|
self.lowercase = lowercase
|
|
|
|
def get_config(self):
|
|
config = {
|
|
'filepath': self.filepath,
|
|
'expand_vocab': self.expand_vocab,
|
|
'lowercase': self.lowercase,
|
|
}
|
|
base_config = super(Word2VecEmbeddingTF, self).get_config()
|
|
base_config['embeddings_initializer'] = self._embeddings_initializer
|
|
return dict(list(base_config.items()) + list(config.items()))
|
|
|
|
|
|
@hanlp_register
|
|
class StringWord2VecEmbeddingTF(Word2VecEmbeddingTF):
|
|
|
|
def __init__(self, filepath: str = None, vocab: VocabTF = None, expand_vocab=True, lowercase=False, input_dim=None,
|
|
output_dim=None, unk=None, normalize=False, embeddings_initializer='VarianceScaling',
|
|
embeddings_regularizer=None, activity_regularizer=None, embeddings_constraint=None, mask_zero=True,
|
|
input_length=None, name=None, **kwargs):
|
|
if vocab is None:
|
|
vocab = VocabTF()
|
|
self.vocab = vocab
|
|
super().__init__(filepath, vocab, expand_vocab, lowercase, input_dim, output_dim, unk, normalize,
|
|
embeddings_initializer, embeddings_regularizer, activity_regularizer, embeddings_constraint,
|
|
mask_zero, input_length, name, **kwargs)
|
|
|
|
def call(self, inputs):
|
|
assert inputs.dtype == tf.string, \
|
|
f'Expect tf.string but got tf.{inputs.dtype.name}. {inputs}' \
|
|
f'Please pass tf.{inputs.dtype.name} in.'
|
|
inputs = self.vocab.lookup(inputs)
|
|
# inputs._keras_mask = tf.not_equal(inputs, self.vocab.pad_idx)
|
|
return super().call(inputs)
|
|
|
|
def compute_mask(self, inputs, mask=None):
|
|
if not self.mask_zero:
|
|
return None
|
|
return tf.not_equal(inputs, self.vocab.pad_token)
|