import { env, ModelRegistry, pipeline } from "@huggingface/transformers";

const modelId = process.argv[2] || "Mike0021/MiniCPM5-1B-ONNX-Web";
const prompt = process.argv[3] || "Hello";

env.allowLocalModels = true;
env.allowRemoteModels = true;

const files = await ModelRegistry.get_pipeline_files("text-generation", modelId, {
  dtype: "q4",
  device: "cpu",
});
console.log("files", JSON.stringify(files, null, 2));

if (!files.includes("onnx/model_q4.onnx")) {
  throw new Error("Expected onnx/model_q4.onnx in Transformers.js file plan.");
}

const generator = await pipeline("text-generation", modelId, {
  dtype: "q4",
  device: "cpu",
});

const result = await generator(prompt, {
  max_new_tokens: 2,
  do_sample: false,
  return_full_text: false,
});

console.log("generation", JSON.stringify(result, null, 2));