# -*- 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)