// SMaLL-100 in transformers.js (@huggingface/transformers v3).
//
//   cd examples/transformers-js && npm install && node translate.mjs
//
// Loads the ONNX + tokenizer.json straight from this repo (local files). For the
// browser, host the repo folder and set env.remoteHost / a model URL instead.
import {
  AutoTokenizer, AutoModelForSeq2SeqLM, Tensor, env,
} from '@huggingface/transformers';
import { readFileSync } from 'node:fs';
import { fileURLToPath } from 'node:url';
import { dirname, resolve } from 'node:path';

const ROOT = resolve(dirname(fileURLToPath(import.meta.url)), '..', '..');
env.allowRemoteModels = false;
env.localModelPath = resolve(ROOT, '..');       // parent dir that contains the model folder
const MODEL = ROOT.split('/').pop();            // this repo's folder name

const LANG = JSON.parse(readFileSync(resolve(ROOT, 'lang_tokens.json'))).lang_to_id;
const tokenizer = await AutoTokenizer.from_pretrained(MODEL);
const model = await AutoModelForSeq2SeqLM.from_pretrained(MODEL);   // int8 onnx (base names)

async function translate(text, tgt) {
  const enc = await tokenizer(text);            // input_ids already end with </s>=2
  const src = Array.from(enc.input_ids.data, Number);
  const ids = [LANG[tgt], ...src];
  const input_ids = new Tensor('int64', BigInt64Array.from(ids.map(BigInt)), [1, ids.length]);
  const attention_mask = new Tensor('int64', BigInt64Array.from(ids.map(() => 1n)), [1, ids.length]);
  const out = await model.generate({ input_ids, attention_mask, max_new_tokens: 128, num_beams: 1 });
  return tokenizer.decode(Array.from(out[0].data, Number), { skip_special_tokens: true });
}

for (const [text, tgt] of [
  ['你好，请问最近的地铁站怎么走？', 'en'],
  ['Excuse me, where can I find a pharmacy?', 'zh'],
  ['この電車は空港に行きますか？', 'ko'],
]) {
  console.log(`[${tgt}] ${text}  ->  ${await translate(text, tgt)}`);
}