Views
No views yet
Fixes an issue with the Qdrant version not having the onnx folder so transformers.js cant use it.
Note: This model is supposed to be used with Qdrant. Vectors have to be configured with Modifier.IDF.
1import { AutoTokenizer, AutoModel, TokenizerModel } from '@xenova/transformers';
2
3documents = [
4 "You should stay, study and sprint.",
5 "History can only prepare us to be surprised yet again.",
6]
7
8const MODEL_ID = "bradynapier/all_miniLM_L6_v2_with_attentions_onnx"
9
10const tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID, {
11 revision: 'main',
12})
13
14// this has some useful utils that transforms py has in the tokenizer ...
15const tokenizerModel = TokenizerModel.fromConfig(tokenizer.model.config)
16
17const model = await AutoModel.from_pretrained(MODEL_ID, {
18 quantized: false,
19 revision: 'main',
20});
21
22
23// the types are wildy incorrect... but this should get you what you need!This may not be the best way but the documentation is truly lacking and this does the job :-P
1/**
2 * Minimal attention tensor shape we rely on.
3 * Only `dims` and `data` are used (dims = [B=1, H, T, T]).
4 */
5type XtTensor = { dims: number[]; data: ArrayLike<number | bigint> };
6
7/**
8 * Collect attentions across layers from a model.forward(...) output.
9 *
10 * ⚠️ Transformers.js variation:
11 * - Some builds return `{ attentions: Tensor[] }`.
12 * - Others return a dict with `attention_1`, `attention_2`, ... per layer.
13 *
14 * @internal
15 * @param out Raw dictionary from `model.forward(...)`.
16 * @returns Array of attention tensors (one per layer) with dims `[1, H, T, T]`.
17 */
18function collectAttentions(out: Record<string, Tensor>): XtTensor[] {
19 // Prefer array form if present (runtime feature; TS types don’t guarantee it).
20 const anyOut = out as unknown as { attentions?: XtTensor[] };
21 if (Array.isArray(anyOut.attentions)) return anyOut.attentions;
22
23 // Otherwise gather attention_1..attention_N and sort numerically by suffix.
24 const keys = Object.keys(out)
25 .filter((k) => /^attention_\d+$/i.test(k))
26 .sort(
27 (a, b) => parseInt(a.split('_')[1], 10) - parseInt(b.split('_')[1], 10),
28 );
29
30 return keys.map((k) => out[k] as unknown as XtTensor);
31}
32
33function onesMask(n: number): Tensor {
34 const data = BigInt64Array.from({ length: n }, () => 1n);
35 return new Tensor('int64', data, [1, n]);
36}
37
38
39/**
40 * Tokenization:
41 * Prefer the public callable form `tokenizer(text, {...})` which returns tensors.
42 * In case your wrapper only exposes a `_call` (private-ish) we fall back to it here.
43 * The return includes `input_ids` and `attention_mask` tensors.
44 */
45const enc =
46 typeof (tokenizer as typeof tokenizer._call) === 'function' ?
47 // eslint-disable-next-line @typescript-eslint/await-thenable
48 await (tokenizer as typeof tokenizer._call)(text, {
49 add_special_tokens: true,
50 })
51 : tokenizer._call(text, { add_special_tokens: true }); // <-- documented hack
52
53// Convert tensor buffers (may be BigInt) → number[] for downstream processing.
54const input_ids = Array.from(
55 (enc.input_ids as Tensor).data as ArrayLike<number | bigint>,
56).map(Number);
57
58/**
59 * Forward pass with attentions.
60 *
61 * Another "crazy" bit: different Transformers.js builds expose attentions differently. We:
62 * - accept `{ attentions: Tensor[] }`, or
63 * - collect `attention_1, attention_2, ...` and sort them.
64 * Also, `Tensor` has no `.get(...)` so we do **flat buffer indexing** with `dims`.
65 */
66const out = (await model.forward({
67 input_ids,
68 attention_mask: onesMask(input_ids.length),
69 output_attentions: true,
70})) as unknown as Record<string, Tensor>;
71
72const attentions = collectAttentions(out)