java实现Word2Vec计算语义相似度例子，详细注释，AI入门学习，附源码，预训练文件

1.准备工作，源码以及预训练文件

源码以及预训练文件比较大下载地址https://pan.quark.cn/s/aeb85eaf95e2

2.核心代码

package org.example;

import com.huaban.analysis.jieba.JiebaSegmenter;
import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.word2vec.Word2Vec;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;
import java.util.stream.Collectors;

public class Main {
    public static void main(String[] args) throws IOException {
        // 输入的两个句子
        String input1 = "一寸光阴一寸金，寸金难买寸光阴。";
        String input2 = "光阴似箭";
        // 词向量模型文件路径
        String modelFile = "baike_26g_news_13g_novel_229g.bin";
        // 读取词向量模型文件
        InputStream is = ClassLoader.getSystemResourceAsStream(modelFile);
        BufferedInputStream bufferedInputStream = new BufferedInputStream(Objects.requireNonNull(is), 1024 * 1024);
        Word2Vec word2VecModel = WordVectorSerializer.readBinaryModel(bufferedInputStream, false, true);
        // 计算并输出两个句子的相似度
        System.out.println(sentenceSimilarity(input1, input2, word2VecModel));
    }

    /**
     * 根据文本内容获取对应的词向量列表
     * @param text 文本内容
     * @param model 词向量模型
     * @return 词向量列表
     */
    private static List<INDArray> getWordVectors(String text, Word2Vec model) {
        // 将文本分词
        List<String> words = segmentWords(text.toLowerCase(Locale.getDefault()));
        // 创建一个列表来存储词向量
        List<INDArray> wordVectors = new ArrayList<>(words.size());
        for (String word : words) {
            if (model.hasWord(word)) {
                wordVectors.add(model.getWordVectorMatrix(word));
            } else {
                // 如果单词不在词汇表中，使用默认向量（这里使用零向量）
                int vectorSize = model.getLayerSize(); // 获取词向量的大小
                INDArray defaultVector = Nd4j.zeros(1, vectorSize); // 创建零向量
                wordVectors.add(defaultVector);
            }
        }
        return wordVectors;
    }

    /**
     * 对句子进行分词处理
     * @param sentence 待分词的句子
     * @return 分词后的词语列表
     */
    private static List<String> segmentWords(String sentence) {
        JiebaSegmenter segmenter = new JiebaSegmenter();
        return segmenter.sentenceProcess(sentence).stream()
                .filter(e -> !" ".equals(e) && !e.isEmpty())
                .collect(Collectors.toList());
    }

    /**
     * 计算两个向量的余弦相似度
     * @param vec1 第一个向量
     * @param vec2 第二个向量
     * @return 余弦相似度值
     */
    private static double cosineSimilarity(INDArray vec1, INDArray vec2) {
        // 计算两个向量的点积
        double dotProduct = vec1.mulRowVector(vec2).sumNumber().doubleValue();
        // 计算两个向量的模长
        double norm1 = vec1.norm2Number().doubleValue();
        double norm2 = vec2.norm2Number().doubleValue();
        // 计算余弦相似度
        return dotProduct / (norm1 * norm2);
    }

    /**
     * 计算两个句子的相似度
     * @param sentence1 第一个句子
     * @param sentence2 第二个句子
     * @param model 词向量模型
     * @return 句子相似度值
     */
    private static double sentenceSimilarity(String sentence1, String sentence2, Word2Vec model) {
        List<INDArray> vectors1 = getWordVectors(sentence1, model);
        List<INDArray> vectors2 = getWordVectors(sentence2, model);
        INDArray avgVector1 = getAverageVector(vectors1, model.getLayerSize());
        INDArray avgVector2 = getAverageVector(vectors2, model.getLayerSize());
        return cosineSimilarity(avgVector1, avgVector2);
    }

    /**
     * 计算一组向量的平均值向量
     * @param vectors 向量列表
     * @param modelSize 向量维度大小
     * @return 平均向量
     */
    private static INDArray getAverageVector(List<INDArray> vectors, int modelSize) {
        INDArray sumVector = Nd4j.zeros(1, modelSize); // 创建一个与第一个向量形状相同的零向量
        for (INDArray vector : vectors) {
            sumVector.addiRowVector(vector); // 使用addi进行原地操作
        }
        INDArray indArray = sumVector.div(vectors.size());
        sumVector.close();
        return indArray; // 将总和除以向量数量以获得平均值
    }
}

4. 预训练文件

4.运行结果

可以下载源码跑一下看看，体验一下语义相似度的魅力源码下载

觉得有用请点赞，有问题请在评论区留言