import { resolve } from 'node:path';
import { AutoTokenizer, PreTrainedModel, env } from '@huggingface/transformers';

if (!process.argv[2]) throw new Error('Usage: node example.mjs /path/to/model');
const directory = resolve(process.argv[2]);
env.allowRemoteModels = false;
const tokenizer = await AutoTokenizer.from_pretrained(directory, { local_files_only: true });
const model = await PreTrainedModel.from_pretrained(directory, {
  dtype: 'q8', local_files_only: true, session_options: { enableCpuMemArena: false }
});
try {
  const query = 'Suggest a vegetarian dinner.';
  const documents = ['The user is vegetarian.', 'The user plays the violin.'];
  for (const document of documents) {
    const inputs = tokenizer([query], { text_pair: [document], padding: true, truncation: false });
    const [batch, length] = inputs.input_ids.dims;
    if (batch !== 1 || length < 1 || length > 8192) throw new Error('Unsupported pair length');
    const { logits } = await model(inputs);
    const score = Number(logits.data[0]);
    if (logits.data.length !== 1 || logits.dims.some(n => n !== 1) || !Number.isFinite(score)) {
      throw new Error('Expected one finite raw logit');
    }
    console.log(JSON.stringify({ document, score }));
  }
} finally {
  await model.dispose();
}