package com.walker.embedding.util;
|
|
import ai.djl.util.Utils;
|
import com.walker.infrastructure.utils.StringUtils;
|
import org.nd4j.linalg.api.ndarray.INDArray;
|
import org.nd4j.linalg.factory.Nd4j;
|
import org.slf4j.Logger;
|
import org.slf4j.LoggerFactory;
|
|
import java.io.File;
|
import java.io.FileInputStream;
|
import java.io.IOException;
|
import java.io.InputStream;
|
import java.nio.file.Path;
|
import java.util.HashMap;
|
import java.util.List;
|
import java.util.Map;
|
|
public final class WordEncoder {
|
|
// private static final String EMPTY_VALUE = "";
|
|
protected final transient Logger logger = LoggerFactory.getLogger(this.getClass());
|
|
// List<String> words = null;
|
private Map<String, Integer> wordsCache = new HashMap<>();
|
private float[][] embeddings = null;
|
|
public WordEncoder(Path vocab, Path embedding) throws IOException {
|
try (InputStream is = new FileInputStream(new File(vocab.toString()))) {
|
List<String> words = Utils.readLines(is, false);
|
int i = 0;
|
for(String w : words){
|
this.wordsCache.put(w, i);
|
i++;
|
}
|
}
|
|
File file = new File(embedding.toString());
|
INDArray array = Nd4j.readNpy(file);
|
embeddings = array.toFloatMatrix();
|
}
|
|
public float[] search(String word) {
|
// for (int i = 0; i < words.size(); i++) {
|
// if (words.get(i).equals(word)) {
|
// return embeddings[i];
|
// }
|
// }
|
if(StringUtils.isEmpty(word)){
|
logger.warn("请提供word单词");
|
return null;
|
}
|
Integer index = this.wordsCache.get(word);
|
if(index == null){
|
logger.debug("未查询到词向量索引:" + word);
|
return null;
|
}
|
return embeddings[index];
|
}
|
|
public float cosineSim(float[] feature1, float[] feature2) {
|
float ret = 0.0f;
|
float mod1 = 0.0f;
|
float mod2 = 0.0f;
|
int length = feature1.length;
|
for (int i = 0; i < length; ++i) {
|
ret += feature1[i] * feature2[i];
|
mod1 += feature1[i] * feature1[i];
|
mod2 += feature2[i] * feature2[i];
|
}
|
return (float) ((ret / Math.sqrt(mod1) / Math.sqrt(mod2) + 1) / 2.0f);
|
}
|
}
|
