package cn.edu.bjut.textmining.chapter3;

import ai.djl.Application;
import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
import ai.djl.inference.Predictor;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ZooModel;
import ai.djl.training.util.ProgressBar;

import java.util.Arrays;

public class BertEmbeddingDemo {
    public static void main(String[] args) {
        try {
            String text1 = "He went to the bank to deposit money.";
            String text2 = "He walked along the river bank.";

            Criteria<String, float[]> criteria = Criteria.builder()
                    .optApplication(Application.NLP.TEXT_EMBEDDING)
                    .setTypes(String.class, float[].class)
                    .optModelUrls("djl://ai.djl.huggingface.pytorch/"
                            + "sentence-transformers/all-MiniLM-L6-v2")
                    .optEngine("PyTorch")
                    .optTranslatorFactory(new TextEmbeddingTranslatorFactory())
                    .optProgress(new ProgressBar())
                    .build();

            try (ZooModel<String, float[]> model = criteria.loadModel();
                 Predictor<String, float[]> predictor = model.newPredictor()) {
                float[] embedding1 = predictor.predict(text1);
                float[] embedding2 = predictor.predict(text2);

                System.out.println("句子1: " + text1);
                System.out.println("向量维度: " + embedding1.length);
                System.out.println("前5位: "
                        + Arrays.toString(Arrays.copyOf(embedding1, 5)));
                System.out.println("句子2: " + text2);
                System.out.println("向量维度: " + embedding2.length);
                System.out.println("前5位: "
                        + Arrays.toString(Arrays.copyOf(embedding2, 5)));
                System.out.printf("余弦相似度: %.4f%n",
                        cosineSimilarity(embedding1, embedding2));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static double cosineSimilarity(float[] vectorA, float[] vectorB) {
        double dotProduct = 0.0;
        double normA = 0.0;
        double normB = 0.0;
        for (int i = 0; i < vectorA.length; i++) {
            dotProduct += vectorA[i] * vectorB[i];
            normA += vectorA[i] * vectorA[i];
            normB += vectorB[i] * vectorB[i];
        }
        return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
    }
}