BIDI_EXTRACTION_AI/BiddingKG/dl/model_runtime/embed.py

# -*- coding: utf-8 -*-
"""词向量 / 字向量加载与 embedding 查找。

按 ARCHITECTURE.md Phase 3 拆分建议，从 ``common/Utils.py`` 迁出。
类型：CORE（模型推理基础，人工主导）。

原位置：``common/Utils.py`` 中以下函数和全局变量：
    - ``model_w2v`` / ``lock_model_w2v`` / ``model_word`` / ``lock_model_word``
    - ``model_word_file``
    - ``getw2vfilepath`` / ``getFileFromSysPath``
    - ``getModel_w2v`` / ``getModel_word``
    - ``embedding`` / ``embedding_word`` / ``embedding_word_forward``
    - ``formEncoding``

本文件自包含，不依赖 ``common/Utils.py``，避免循环 import。
``common/Utils.py`` 仍 re-export 以上全部名称，老 import 不受影响。

按 ARCHITECTURE.md §4.3 依赖方向约束：
    model_runtime -> domain, infra only
"""
from __future__ import absolute_import

import os
import sys
import re
from threading import RLock

import numpy as np
import gensim

__all__ = [
    "model_w2v",
    "lock_model_w2v",
    "model_word",
    "lock_model_word",
    "model_word_file",
    "getw2vfilepath",
    "getFileFromSysPath",
    "getModel_w2v",
    "getModel_word",
    "embedding",
    "embedding_word",
    "embedding_word_forward",
    "formEncoding",
]

model_w2v = None
lock_model_w2v = RLock()

model_word_file = os.path.dirname(os.path.abspath(__file__)) + "/../singlew2v_model.vector"
model_word = None
lock_model_word = RLock()

Lazy_load = False


def getLazyLoad():
    global Lazy_load
    return Lazy_load


def getFileFromSysPath(filename):
    for _path in sys.path:
        if os.path.isdir(_path):
            for _file in os.listdir(_path):
                _abspath = os.path.join(_path, _file)
                if os.path.isfile(_abspath):
                    if _file == filename:
                        return _abspath
    return None


def getw2vfilepath():
    filename = "wiki_128_word_embedding_new.vector"
    w2vfile = getFileFromSysPath(filename)
    if w2vfile is not None:
        return w2vfile
    return filename


def getModel_w2v():
    '''
    @summary:加载词向量
    '''
    global model_w2v, lock_model_w2v
    with lock_model_w2v:
        if model_w2v is None:
            model_w2v = gensim.models.KeyedVectors.load_word2vec_format(getw2vfilepath(), binary=True)
        return model_w2v


def getModel_word():
    '''
    @summary:加载字向量
    '''

    global model_word, lock_model_w2v
    with lock_model_word:
        if model_word is None:
            model_word = gensim.models.KeyedVectors.load_word2vec_format(model_word_file, binary=True)
        return model_word


def embedding(datas, shape):
    '''
    @summary:查找词汇对应的词向量
    @param:
        datas:词汇的list
        shape:结果的shape
    @return: array,返回对应shape的词嵌入
    '''
    model_w2v = getModel_w2v()
    embed = np.zeros(shape)
    length = shape[1]
    out_index = 0
    for data in datas:
        index = 0
        for item in data:
            item_not_space = re.sub("\s*", "", item)
            if index >= length:
                break
            if item_not_space in model_w2v.vocab:
                embed[out_index][index] = model_w2v[item_not_space]
                index += 1
            else:
                index += 1
        out_index += 1
    return embed


def embedding_word(datas, shape):
    '''
    @summary:查找词汇对应的词向量
    @param:
        datas:词汇的list
        shape:结果的shape
    @return: array,返回对应shape的词嵌入
    '''
    model_w2v = getModel_word()
    embed = np.zeros(shape)
    length = shape[1]
    out_index = 0
    for data in datas:
        index = 0
        for item in str(data)[-shape[1]:]:
            if index >= length:
                break
            if item in model_w2v.vocab:
                embed[out_index][index] = model_w2v[item]
                index += 1
            else:
                index += 1
        out_index += 1
    return embed


def embedding_word_forward(datas, shape):
    '''
    @summary:查找词汇对应的词向量
    @param:
        datas:词汇的list
        shape:结果的shape
    @return: array,返回对应shape的词嵌入
    '''
    model_w2v = getModel_word()
    embed = np.zeros(shape)
    length = shape[1]
    out_index = 0
    for data in datas:
        index = 0
        for item in str(data)[:shape[1]]:
            if index >= length:
                break
            if item in model_w2v.vocab:
                embed[out_index][index] = model_w2v[item]
                index += 1
            else:
                index += 1
        out_index += 1
    return embed


def formEncoding(text, shape=(100, 60), expand=False):
    embedding = np.zeros(shape)
    word_model = getModel_word()
    for i in range(len(text)):
        if i >= shape[0]:
            break
        if text[i] in word_model.vocab:
            embedding[i] = word_model[text[i]]
    if expand:
        embedding = np.expand_dims(embedding, 0)
    return embedding