| 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(); |
| } |
|
|
