/*
 * This Java source file was generated by the Gradle 'init' task.
 */
package demo4j;

import java.util.Arrays;
import java.util.Map;
import java.nio.file.Paths;
import ai.onnxruntime.*;
import ai.onnxruntime.extensions.OrtxPackage;

public class App {
    public String inference(){
        try {
            var env = OrtEnvironment.getEnvironment();
            var sess_opt = new OrtSession.SessionOptions();

            /* Register the custom ops from onnxruntime-extensions */
            sess_opt.registerCustomOpLibrary(OrtxPackage.getLibraryPath());

            /* do a quick inference on Bert Tokenizer custom ops with Ort */
            var modelPath = Paths.get(this.getClass().getClassLoader().getResource("test_bert_tokenizer.onnx").getPath());
            var session = env.createSession(modelPath.toString(), sess_opt);
            var t1 = OnnxTensor.createTensor(env, new String[]{"This is a test"});
            var inputs = Map.of("text", t1);
            try (var r = session.run(inputs)) {
                long[] tokenIds = (long[])r.get("input_ids").get().getValue();
                return Arrays.toString(tokenIds);
            }

        } catch(OrtException e1) {
            return e1.getMessage();
        }
    }

    public static void main(String[] args) {
        System.out.println(new App().inference());
    }
}
