package com.sd.hmct_small_model;

import org.pytorch.IValue;
import org.pytorch.Module;
import org.pytorch.Tensor;

import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class PyTvmJava {
    int TOP_K = 3;

    /** Executor to run inference task in the background. */
    private ExecutorService executorService  = Executors.newCachedThreadPool();

    public void analyzeText(String text, Module mModule) {
        executorService.execute(() -> {
            try {
//                val bytes: ByteArray = text.encodeToByteArray()
                byte[] bytes = text.getBytes();
//                val shape = longArrayOf(1, bytes.size.toLong())
                long[] shape = new long[]{1L, bytes.length};
//                val inputTensor = Tensor.fromBlobUnsigned(bytes, shape)
                Tensor inputTensor = Tensor.fromBlobUnsigned(bytes, shape);
//                val outputTensor: Tensor = mModule!!.forward(IValue.from(inputTensor)).toTensor()
                Tensor outputTensor = mModule.forward(IValue.from(inputTensor)).toTensor();
//                val scores = outputTensor.getDataAsFloatArray();
                float[] scores = outputTensor.getDataAsFloatArray();
//                val ixs: IntArray? = Utils.topK(scores, TOP_K)
                int[] ixs = Utils.topK(scores, TOP_K);
            } catch (Exception e) {
                e.printStackTrace();
            }
        });

    }

}
