qwen-embedder 0.1.1 → 0.1.4
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/index.d.ts +12 -1
- package/dist/index.js +49 -11
- package/dist/index.js.map +1 -1
- package/dist/sync-worker.d.ts +2 -0
- package/dist/sync-worker.js +81 -0
- package/dist/sync-worker.js.map +1 -0
- package/package.json +6 -4
- package/src/embedder.ts +28 -10
- package/src/index.ts +2 -0
- package/src/sync-embedder.ts +34 -0
- package/src/sync-worker.ts +96 -0
package/dist/index.d.ts
CHANGED
|
@@ -4,12 +4,23 @@ interface EmbedderOptions {
|
|
|
4
4
|
cacheSize?: number;
|
|
5
5
|
}
|
|
6
6
|
|
|
7
|
+
interface SyncEmbedderOptions {
|
|
8
|
+
cacheSize?: number;
|
|
9
|
+
}
|
|
10
|
+
declare class SyncEmbedder {
|
|
11
|
+
private fn;
|
|
12
|
+
constructor(options?: SyncEmbedderOptions);
|
|
13
|
+
embedSync(text: string): number[];
|
|
14
|
+
embedSync(texts: string[]): number[][];
|
|
15
|
+
}
|
|
16
|
+
|
|
7
17
|
declare class Embedder {
|
|
8
18
|
private pipe;
|
|
9
19
|
private cache;
|
|
10
20
|
private maxCacheSize;
|
|
11
21
|
private queue;
|
|
12
22
|
private constructor();
|
|
23
|
+
static createSync(options?: SyncEmbedderOptions): SyncEmbedder;
|
|
13
24
|
static create(options?: EmbedderOptions): Promise<Embedder>;
|
|
14
25
|
embed(text: string): Promise<number[]>;
|
|
15
26
|
embed(texts: string[]): Promise<number[][]>;
|
|
@@ -19,4 +30,4 @@ declare class Embedder {
|
|
|
19
30
|
private setCache;
|
|
20
31
|
}
|
|
21
32
|
|
|
22
|
-
export { Embedder, type EmbedderOptions };
|
|
33
|
+
export { Embedder, type EmbedderOptions, SyncEmbedder, type SyncEmbedderOptions };
|
package/dist/index.js
CHANGED
|
@@ -3,6 +3,28 @@ import { pipeline } from "@huggingface/transformers";
|
|
|
3
3
|
import PQueue from "p-queue";
|
|
4
4
|
import { homedir } from "os";
|
|
5
5
|
import { join } from "path";
|
|
6
|
+
|
|
7
|
+
// src/sync-embedder.ts
|
|
8
|
+
import { createSyncFn } from "synckit";
|
|
9
|
+
import { fileURLToPath } from "url";
|
|
10
|
+
import { dirname, resolve } from "path";
|
|
11
|
+
var __dirname = dirname(fileURLToPath(import.meta.url));
|
|
12
|
+
var workerPath = resolve(__dirname, "sync-worker.js");
|
|
13
|
+
var SyncEmbedder = class {
|
|
14
|
+
fn;
|
|
15
|
+
constructor(options = {}) {
|
|
16
|
+
this.fn = createSyncFn(workerPath);
|
|
17
|
+
this.fn({ _cmd: "init", cacheSize: options.cacheSize ?? 100 });
|
|
18
|
+
}
|
|
19
|
+
embedSync(input) {
|
|
20
|
+
if (Array.isArray(input)) {
|
|
21
|
+
return this.fn({ _cmd: "embed", texts: input });
|
|
22
|
+
}
|
|
23
|
+
return this.fn({ _cmd: "embed", text: input });
|
|
24
|
+
}
|
|
25
|
+
};
|
|
26
|
+
|
|
27
|
+
// src/embedder.ts
|
|
6
28
|
var MODEL_ID = "onnx-community/Qwen3-Embedding-0.6B-ONNX";
|
|
7
29
|
var CACHE_DIR = join(homedir(), ".qwenembedder", ".cache", "models");
|
|
8
30
|
function resolveConcurrency(options) {
|
|
@@ -21,6 +43,9 @@ var Embedder = class _Embedder {
|
|
|
21
43
|
this.cache = /* @__PURE__ */ new Map();
|
|
22
44
|
this.queue = concurrency !== null ? new PQueue({ concurrency }) : null;
|
|
23
45
|
}
|
|
46
|
+
static createSync(options = {}) {
|
|
47
|
+
return new SyncEmbedder(options);
|
|
48
|
+
}
|
|
24
49
|
static async create(options = {}) {
|
|
25
50
|
const cacheSize = options.cacheSize ?? 100;
|
|
26
51
|
const concurrency = resolveConcurrency(options);
|
|
@@ -53,7 +78,8 @@ var Embedder = class _Embedder {
|
|
|
53
78
|
return [];
|
|
54
79
|
}
|
|
55
80
|
const results = new Array(texts.length);
|
|
56
|
-
const
|
|
81
|
+
const uncachedByText = /* @__PURE__ */ new Map();
|
|
82
|
+
const uncachedTexts = [];
|
|
57
83
|
for (let i = 0; i < texts.length; i++) {
|
|
58
84
|
const t = texts[i];
|
|
59
85
|
if (typeof t !== "string" || t.length === 0) {
|
|
@@ -61,25 +87,35 @@ var Embedder = class _Embedder {
|
|
|
61
87
|
}
|
|
62
88
|
const cached = this.cache.get(t);
|
|
63
89
|
if (cached !== void 0) {
|
|
64
|
-
results[i] = cached;
|
|
90
|
+
results[i] = cached.slice();
|
|
65
91
|
} else {
|
|
66
|
-
|
|
92
|
+
const existing = uncachedByText.get(t);
|
|
93
|
+
if (existing !== void 0) {
|
|
94
|
+
existing.push(i);
|
|
95
|
+
} else {
|
|
96
|
+
uncachedByText.set(t, [i]);
|
|
97
|
+
uncachedTexts.push(t);
|
|
98
|
+
}
|
|
67
99
|
}
|
|
68
100
|
}
|
|
69
|
-
if (
|
|
101
|
+
if (uncachedTexts.length === 0) {
|
|
70
102
|
return results;
|
|
71
103
|
}
|
|
72
|
-
const uncachedTexts = uncached.map((i) => texts[i]);
|
|
73
104
|
const inferred = await this.runInference(uncachedTexts);
|
|
74
|
-
for (let j = 0; j <
|
|
105
|
+
for (let j = 0; j < uncachedTexts.length; j++) {
|
|
106
|
+
const text = uncachedTexts[j];
|
|
75
107
|
const vector = inferred[j];
|
|
76
|
-
|
|
77
|
-
|
|
108
|
+
const stored = vector.slice();
|
|
109
|
+
this.setCache(text, stored);
|
|
110
|
+
const positions = uncachedByText.get(text) ?? [];
|
|
111
|
+
for (const position of positions) {
|
|
112
|
+
results[position] = stored.slice();
|
|
113
|
+
}
|
|
78
114
|
}
|
|
79
115
|
return results;
|
|
80
116
|
}
|
|
81
117
|
async runInference(input) {
|
|
82
|
-
const exec = () => this.pipe(input, { pooling: "
|
|
118
|
+
const exec = () => this.pipe(input, { pooling: "last_token", normalize: true });
|
|
83
119
|
const output = this.queue ? await this.queue.add(exec) : await exec();
|
|
84
120
|
const data = Array.from(output.data);
|
|
85
121
|
const hiddenDim = output.dims[output.dims.length - 1];
|
|
@@ -94,16 +130,18 @@ var Embedder = class _Embedder {
|
|
|
94
130
|
}
|
|
95
131
|
setCache(key, value) {
|
|
96
132
|
if (this.maxCacheSize === 0) return;
|
|
133
|
+
const stored = value.slice();
|
|
97
134
|
if (this.cache.size >= this.maxCacheSize) {
|
|
98
135
|
const oldest = this.cache.keys().next();
|
|
99
136
|
if (!oldest.done) {
|
|
100
137
|
this.cache.delete(oldest.value);
|
|
101
138
|
}
|
|
102
139
|
}
|
|
103
|
-
this.cache.set(key,
|
|
140
|
+
this.cache.set(key, stored);
|
|
104
141
|
}
|
|
105
142
|
};
|
|
106
143
|
export {
|
|
107
|
-
Embedder
|
|
144
|
+
Embedder,
|
|
145
|
+
SyncEmbedder
|
|
108
146
|
};
|
|
109
147
|
//# sourceMappingURL=index.js.map
|
package/dist/index.js.map
CHANGED
|
@@ -1 +1 @@
|
|
|
1
|
-
{"version":3,"sources":["../src/embedder.ts"],"sourcesContent":["import { pipeline } from '@huggingface/transformers'\nimport PQueue from 'p-queue'\nimport { homedir } from 'os'\nimport { join } from 'path'\nimport type { EmbedderOptions } from './types.js'\n\nconst MODEL_ID = 'onnx-community/Qwen3-Embedding-0.6B-ONNX'\nconst CACHE_DIR = join(homedir(), '.qwenembedder', '.cache', 'models')\n\nfunction resolveConcurrency(options: EmbedderOptions): number | null {\n if (options.concurrency !== undefined) return options.concurrency\n if (options.queue) return 1\n return null\n}\n\ninterface InferenceResult {\n data: Float32Array | number[]\n dims: number[]\n tolist(): number[][]\n}\n\nexport class Embedder {\n private pipe: (texts: string | string[], options?: Record<string, unknown>) => Promise<InferenceResult>\n private cache: Map<string, number[]>\n private maxCacheSize: number\n private queue: PQueue | null\n\n private constructor(\n pipe: Embedder['pipe'],\n concurrency: number | null,\n maxCacheSize: number,\n ) {\n this.pipe = pipe\n this.maxCacheSize = maxCacheSize\n this.cache = new Map()\n this.queue = concurrency !== null ? new PQueue({ concurrency }) : null\n }\n\n static async create(options: EmbedderOptions = {}): Promise<Embedder> {\n const cacheSize = options.cacheSize ?? 100\n const concurrency = resolveConcurrency(options)\n\n const pipe = await pipeline('feature-extraction', MODEL_ID, {\n cache_dir: CACHE_DIR,\n dtype: 'q8',\n })\n\n return new Embedder(pipe as unknown as Embedder['pipe'], concurrency, cacheSize)\n }\n\n async embed(text: string): Promise<number[]>\n async embed(texts: string[]): Promise<number[][]>\n async embed(input: string | string[]): Promise<number[] | number[][]> {\n if (Array.isArray(input)) {\n return this.embedBatch(input)\n }\n return this.embedSingle(input)\n }\n\n private async embedSingle(text: string): Promise<number[]> {\n if (text.length === 0) {\n throw new TypeError('Input must be a non-empty string')\n }\n\n const cached = this.cache.get(text)\n if (cached !== undefined) {\n return cached.slice()\n }\n\n const result = await this.runInference(text)\n this.setCache(text, result)\n return result.slice()\n }\n\n private async embedBatch(texts: string[]): Promise<number[][]> {\n if (texts.length === 0) {\n return []\n }\n\n const results: (number[] | undefined)[] = new Array(texts.length)\n const
|
|
1
|
+
{"version":3,"sources":["../src/embedder.ts","../src/sync-embedder.ts"],"sourcesContent":["import { pipeline } from '@huggingface/transformers'\nimport PQueue from 'p-queue'\nimport { homedir } from 'os'\nimport { join } from 'path'\nimport type { EmbedderOptions } from './types.js'\nimport { SyncEmbedder } from './sync-embedder.js'\nimport type { SyncEmbedderOptions } from './sync-embedder.js'\n\nconst MODEL_ID = 'onnx-community/Qwen3-Embedding-0.6B-ONNX'\nconst CACHE_DIR = join(homedir(), '.qwenembedder', '.cache', 'models')\n\nfunction resolveConcurrency(options: EmbedderOptions): number | null {\n if (options.concurrency !== undefined) return options.concurrency\n if (options.queue) return 1\n return null\n}\n\ninterface InferenceResult {\n data: Float32Array | number[]\n dims: number[]\n tolist(): number[][]\n}\n\nexport class Embedder {\n private pipe: (texts: string | string[], options?: Record<string, unknown>) => Promise<InferenceResult>\n private cache: Map<string, number[]>\n private maxCacheSize: number\n private queue: PQueue | null\n\n private constructor(\n pipe: Embedder['pipe'],\n concurrency: number | null,\n maxCacheSize: number,\n ) {\n this.pipe = pipe\n this.maxCacheSize = maxCacheSize\n this.cache = new Map()\n this.queue = concurrency !== null ? new PQueue({ concurrency }) : null\n }\n\n static createSync(options: SyncEmbedderOptions = {}): SyncEmbedder {\n return new SyncEmbedder(options)\n }\n\n static async create(options: EmbedderOptions = {}): Promise<Embedder> {\n const cacheSize = options.cacheSize ?? 100\n const concurrency = resolveConcurrency(options)\n\n const pipe = await pipeline('feature-extraction', MODEL_ID, {\n cache_dir: CACHE_DIR,\n dtype: 'q8',\n })\n\n return new Embedder(pipe as unknown as Embedder['pipe'], concurrency, cacheSize)\n }\n\n async embed(text: string): Promise<number[]>\n async embed(texts: string[]): Promise<number[][]>\n async embed(input: string | string[]): Promise<number[] | number[][]> {\n if (Array.isArray(input)) {\n return this.embedBatch(input)\n }\n return this.embedSingle(input)\n }\n\n private async embedSingle(text: string): Promise<number[]> {\n if (text.length === 0) {\n throw new TypeError('Input must be a non-empty string')\n }\n\n const cached = this.cache.get(text)\n if (cached !== undefined) {\n return cached.slice()\n }\n\n const result = await this.runInference(text)\n this.setCache(text, result)\n return result.slice()\n }\n\n private async embedBatch(texts: string[]): Promise<number[][]> {\n if (texts.length === 0) {\n return []\n }\n\n const results: (number[] | undefined)[] = new Array(texts.length)\n const uncachedByText = new Map<string, number[]>()\n const uncachedTexts: string[] = []\n\n for (let i = 0; i < texts.length; i++) {\n const t = texts[i]\n if (typeof t !== 'string' || t.length === 0) {\n throw new TypeError('Input must be a non-empty string')\n }\n const cached = this.cache.get(t)\n if (cached !== undefined) {\n results[i] = cached.slice()\n } else {\n const existing = uncachedByText.get(t)\n if (existing !== undefined) {\n existing.push(i)\n } else {\n uncachedByText.set(t, [i])\n uncachedTexts.push(t)\n }\n }\n }\n\n if (uncachedTexts.length === 0) {\n return results as number[][]\n }\n\n const inferred = await this.runInference(uncachedTexts)\n\n for (let j = 0; j < uncachedTexts.length; j++) {\n const text = uncachedTexts[j]\n const vector = inferred[j]\n const stored = vector.slice()\n this.setCache(text, stored)\n const positions = uncachedByText.get(text) ?? []\n for (const position of positions) {\n results[position] = stored.slice()\n }\n }\n\n return results as number[][]\n }\n\n private async runInference(text: string): Promise<number[]>\n private async runInference(texts: string[]): Promise<number[][]>\n private async runInference(input: string | string[]): Promise<number[] | number[][]> {\n const exec = () => this.pipe(input, { pooling: 'last_token', normalize: true })\n\n const output = this.queue\n ? await this.queue.add(exec)\n : await exec()\n\n const data = Array.from(output.data) as number[]\n const hiddenDim = output.dims[output.dims.length - 1]\n\n if (typeof input === 'string') {\n return data\n }\n\n const result: number[][] = []\n for (let i = 0; i < data.length; i += hiddenDim) {\n result.push(data.slice(i, i + hiddenDim))\n }\n return result\n }\n\n private setCache(key: string, value: number[]): void {\n if (this.maxCacheSize === 0) return\n\n const stored = value.slice()\n if (this.cache.size >= this.maxCacheSize) {\n const oldest = this.cache.keys().next()\n if (!oldest.done) {\n this.cache.delete(oldest.value)\n }\n }\n this.cache.set(key, stored)\n }\n}\n","import { createSyncFn } from 'synckit'\nimport { fileURLToPath } from 'url'\nimport { dirname, resolve } from 'path'\n\nconst __dirname = dirname(fileURLToPath(import.meta.url))\nconst workerPath = resolve(__dirname, 'sync-worker.js')\n\ntype SyncFn = {\n (command: { _cmd: 'init'; cacheSize?: number }): void\n (command: { _cmd: 'embed'; text: string }): number[]\n (command: { _cmd: 'embed'; texts: string[] }): number[][]\n}\n\nexport interface SyncEmbedderOptions {\n cacheSize?: number\n}\n\nexport class SyncEmbedder {\n private fn: SyncFn\n\n constructor(options: SyncEmbedderOptions = {}) {\n this.fn = createSyncFn(workerPath) as SyncFn\n this.fn({ _cmd: 'init', cacheSize: options.cacheSize ?? 100 })\n }\n\n embedSync(text: string): number[]\n embedSync(texts: string[]): number[][]\n embedSync(input: string | string[]): number[] | number[][] {\n if (Array.isArray(input)) {\n return this.fn({ _cmd: 'embed', texts: input })\n }\n return this.fn({ _cmd: 'embed', text: input })\n }\n}\n"],"mappings":";AAAA,SAAS,gBAAgB;AACzB,OAAO,YAAY;AACnB,SAAS,eAAe;AACxB,SAAS,YAAY;;;ACHrB,SAAS,oBAAoB;AAC7B,SAAS,qBAAqB;AAC9B,SAAS,SAAS,eAAe;AAEjC,IAAM,YAAY,QAAQ,cAAc,YAAY,GAAG,CAAC;AACxD,IAAM,aAAa,QAAQ,WAAW,gBAAgB;AAY/C,IAAM,eAAN,MAAmB;AAAA,EAChB;AAAA,EAER,YAAY,UAA+B,CAAC,GAAG;AAC7C,SAAK,KAAK,aAAa,UAAU;AACjC,SAAK,GAAG,EAAE,MAAM,QAAQ,WAAW,QAAQ,aAAa,IAAI,CAAC;AAAA,EAC/D;AAAA,EAIA,UAAU,OAAiD;AACzD,QAAI,MAAM,QAAQ,KAAK,GAAG;AACxB,aAAO,KAAK,GAAG,EAAE,MAAM,SAAS,OAAO,MAAM,CAAC;AAAA,IAChD;AACA,WAAO,KAAK,GAAG,EAAE,MAAM,SAAS,MAAM,MAAM,CAAC;AAAA,EAC/C;AACF;;;ADzBA,IAAM,WAAW;AACjB,IAAM,YAAY,KAAK,QAAQ,GAAG,iBAAiB,UAAU,QAAQ;AAErE,SAAS,mBAAmB,SAAyC;AACnE,MAAI,QAAQ,gBAAgB,OAAW,QAAO,QAAQ;AACtD,MAAI,QAAQ,MAAO,QAAO;AAC1B,SAAO;AACT;AAQO,IAAM,WAAN,MAAM,UAAS;AAAA,EACZ;AAAA,EACA;AAAA,EACA;AAAA,EACA;AAAA,EAEA,YACN,MACA,aACA,cACA;AACA,SAAK,OAAO;AACZ,SAAK,eAAe;AACpB,SAAK,QAAQ,oBAAI,IAAI;AACrB,SAAK,QAAQ,gBAAgB,OAAO,IAAI,OAAO,EAAE,YAAY,CAAC,IAAI;AAAA,EACpE;AAAA,EAEA,OAAO,WAAW,UAA+B,CAAC,GAAiB;AACjE,WAAO,IAAI,aAAa,OAAO;AAAA,EACjC;AAAA,EAEA,aAAa,OAAO,UAA2B,CAAC,GAAsB;AACpE,UAAM,YAAY,QAAQ,aAAa;AACvC,UAAM,cAAc,mBAAmB,OAAO;AAE9C,UAAM,OAAO,MAAM,SAAS,sBAAsB,UAAU;AAAA,MAC1D,WAAW;AAAA,MACX,OAAO;AAAA,IACT,CAAC;AAED,WAAO,IAAI,UAAS,MAAqC,aAAa,SAAS;AAAA,EACjF;AAAA,EAIA,MAAM,MAAM,OAA0D;AACpE,QAAI,MAAM,QAAQ,KAAK,GAAG;AACxB,aAAO,KAAK,WAAW,KAAK;AAAA,IAC9B;AACA,WAAO,KAAK,YAAY,KAAK;AAAA,EAC/B;AAAA,EAEA,MAAc,YAAY,MAAiC;AACzD,QAAI,KAAK,WAAW,GAAG;AACrB,YAAM,IAAI,UAAU,kCAAkC;AAAA,IACxD;AAEA,UAAM,SAAS,KAAK,MAAM,IAAI,IAAI;AAClC,QAAI,WAAW,QAAW;AACxB,aAAO,OAAO,MAAM;AAAA,IACtB;AAEA,UAAM,SAAS,MAAM,KAAK,aAAa,IAAI;AAC3C,SAAK,SAAS,MAAM,MAAM;AAC1B,WAAO,OAAO,MAAM;AAAA,EACtB;AAAA,EAEA,MAAc,WAAW,OAAsC;AAC7D,QAAI,MAAM,WAAW,GAAG;AACtB,aAAO,CAAC;AAAA,IACV;AAEA,UAAM,UAAoC,IAAI,MAAM,MAAM,MAAM;AAChE,UAAM,iBAAiB,oBAAI,IAAsB;AACjD,UAAM,gBAA0B,CAAC;AAEjC,aAAS,IAAI,GAAG,IAAI,MAAM,QAAQ,KAAK;AACrC,YAAM,IAAI,MAAM,CAAC;AACjB,UAAI,OAAO,MAAM,YAAY,EAAE,WAAW,GAAG;AAC3C,cAAM,IAAI,UAAU,kCAAkC;AAAA,MACxD;AACA,YAAM,SAAS,KAAK,MAAM,IAAI,CAAC;AAC/B,UAAI,WAAW,QAAW;AACxB,gBAAQ,CAAC,IAAI,OAAO,MAAM;AAAA,MAC5B,OAAO;AACL,cAAM,WAAW,eAAe,IAAI,CAAC;AACrC,YAAI,aAAa,QAAW;AAC1B,mBAAS,KAAK,CAAC;AAAA,QACjB,OAAO;AACL,yBAAe,IAAI,GAAG,CAAC,CAAC,CAAC;AACzB,wBAAc,KAAK,CAAC;AAAA,QACtB;AAAA,MACF;AAAA,IACF;AAEA,QAAI,cAAc,WAAW,GAAG;AAC9B,aAAO;AAAA,IACT;AAEA,UAAM,WAAW,MAAM,KAAK,aAAa,aAAa;AAEtD,aAAS,IAAI,GAAG,IAAI,cAAc,QAAQ,KAAK;AAC7C,YAAM,OAAO,cAAc,CAAC;AAC5B,YAAM,SAAS,SAAS,CAAC;AACzB,YAAM,SAAS,OAAO,MAAM;AAC5B,WAAK,SAAS,MAAM,MAAM;AAC1B,YAAM,YAAY,eAAe,IAAI,IAAI,KAAK,CAAC;AAC/C,iBAAW,YAAY,WAAW;AAChC,gBAAQ,QAAQ,IAAI,OAAO,MAAM;AAAA,MACnC;AAAA,IACF;AAEA,WAAO;AAAA,EACT;AAAA,EAIA,MAAc,aAAa,OAA0D;AACnF,UAAM,OAAO,MAAM,KAAK,KAAK,OAAO,EAAE,SAAS,cAAc,WAAW,KAAK,CAAC;AAE9E,UAAM,SAAS,KAAK,QAChB,MAAM,KAAK,MAAM,IAAI,IAAI,IACzB,MAAM,KAAK;AAEf,UAAM,OAAO,MAAM,KAAK,OAAO,IAAI;AACnC,UAAM,YAAY,OAAO,KAAK,OAAO,KAAK,SAAS,CAAC;AAEpD,QAAI,OAAO,UAAU,UAAU;AAC7B,aAAO;AAAA,IACT;AAEA,UAAM,SAAqB,CAAC;AAC5B,aAAS,IAAI,GAAG,IAAI,KAAK,QAAQ,KAAK,WAAW;AAC/C,aAAO,KAAK,KAAK,MAAM,GAAG,IAAI,SAAS,CAAC;AAAA,IAC1C;AACA,WAAO;AAAA,EACT;AAAA,EAEQ,SAAS,KAAa,OAAuB;AACnD,QAAI,KAAK,iBAAiB,EAAG;AAE7B,UAAM,SAAS,MAAM,MAAM;AAC3B,QAAI,KAAK,MAAM,QAAQ,KAAK,cAAc;AACxC,YAAM,SAAS,KAAK,MAAM,KAAK,EAAE,KAAK;AACtC,UAAI,CAAC,OAAO,MAAM;AAChB,aAAK,MAAM,OAAO,OAAO,KAAK;AAAA,MAChC;AAAA,IACF;AACA,SAAK,MAAM,IAAI,KAAK,MAAM;AAAA,EAC5B;AACF;","names":[]}
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
// src/sync-worker.ts
|
|
2
|
+
import { pipeline } from "@huggingface/transformers";
|
|
3
|
+
import { runAsWorker } from "synckit";
|
|
4
|
+
import { homedir } from "os";
|
|
5
|
+
import { join } from "path";
|
|
6
|
+
var MODEL_ID = "onnx-community/Qwen3-Embedding-0.6B-ONNX";
|
|
7
|
+
var CACHE_DIR = join(homedir(), ".qwenembedder", ".cache", "models");
|
|
8
|
+
var pipe = null;
|
|
9
|
+
var cache = null;
|
|
10
|
+
var maxCacheSize = 100;
|
|
11
|
+
function setCache(key, value) {
|
|
12
|
+
if (!cache) return;
|
|
13
|
+
if (maxCacheSize === 0) return;
|
|
14
|
+
if (cache.size >= maxCacheSize) {
|
|
15
|
+
const oldest = cache.keys().next();
|
|
16
|
+
if (!oldest.done) {
|
|
17
|
+
cache.delete(oldest.value);
|
|
18
|
+
}
|
|
19
|
+
}
|
|
20
|
+
cache.set(key, value);
|
|
21
|
+
}
|
|
22
|
+
runAsWorker(async (command) => {
|
|
23
|
+
if (command._cmd === "init") {
|
|
24
|
+
maxCacheSize = command.cacheSize ?? 100;
|
|
25
|
+
cache = /* @__PURE__ */ new Map();
|
|
26
|
+
return;
|
|
27
|
+
}
|
|
28
|
+
if (!pipe) {
|
|
29
|
+
pipe = await pipeline("feature-extraction", MODEL_ID, {
|
|
30
|
+
cache_dir: CACHE_DIR,
|
|
31
|
+
dtype: "q8"
|
|
32
|
+
});
|
|
33
|
+
}
|
|
34
|
+
if (command._cmd === "embed") {
|
|
35
|
+
const texts = command.texts ?? (command.text ? [command.text] : []);
|
|
36
|
+
if (texts.length === 0) {
|
|
37
|
+
throw new TypeError("Input must be a non-empty string");
|
|
38
|
+
}
|
|
39
|
+
const results = new Array(texts.length);
|
|
40
|
+
const uncachedByText = /* @__PURE__ */ new Map();
|
|
41
|
+
const uncachedTexts = [];
|
|
42
|
+
for (let i = 0; i < texts.length; i++) {
|
|
43
|
+
const t = texts[i];
|
|
44
|
+
if (typeof t !== "string" || t.length === 0) {
|
|
45
|
+
throw new TypeError("Input must be a non-empty string");
|
|
46
|
+
}
|
|
47
|
+
const cached = cache?.get(t);
|
|
48
|
+
if (cached !== void 0) {
|
|
49
|
+
results[i] = cached.slice();
|
|
50
|
+
} else {
|
|
51
|
+
const existing = uncachedByText.get(t);
|
|
52
|
+
if (existing !== void 0) {
|
|
53
|
+
existing.push(i);
|
|
54
|
+
} else {
|
|
55
|
+
uncachedByText.set(t, [i]);
|
|
56
|
+
uncachedTexts.push(t);
|
|
57
|
+
}
|
|
58
|
+
}
|
|
59
|
+
}
|
|
60
|
+
if (uncachedTexts.length > 0) {
|
|
61
|
+
const output = await pipe(uncachedTexts, { pooling: "last_token", normalize: true });
|
|
62
|
+
const data = Array.from(output.data);
|
|
63
|
+
const hiddenDim = output.dims[output.dims.length - 1];
|
|
64
|
+
for (let j = 0; j < uncachedTexts.length; j++) {
|
|
65
|
+
const text = uncachedTexts[j];
|
|
66
|
+
const vector = data.slice(j * hiddenDim, (j + 1) * hiddenDim);
|
|
67
|
+
const stored = vector.slice();
|
|
68
|
+
setCache(text, stored);
|
|
69
|
+
const positions = uncachedByText.get(text) ?? [];
|
|
70
|
+
for (const position of positions) {
|
|
71
|
+
results[position] = stored.slice();
|
|
72
|
+
}
|
|
73
|
+
}
|
|
74
|
+
}
|
|
75
|
+
if (command.text !== void 0) {
|
|
76
|
+
return results[0];
|
|
77
|
+
}
|
|
78
|
+
return results;
|
|
79
|
+
}
|
|
80
|
+
});
|
|
81
|
+
//# sourceMappingURL=sync-worker.js.map
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
{"version":3,"sources":["../src/sync-worker.ts"],"sourcesContent":["import { pipeline } from '@huggingface/transformers'\nimport { runAsWorker } from 'synckit'\nimport { homedir } from 'os'\nimport { join } from 'path'\n\ntype PipeFn = (texts: string[], options?: Record<string, unknown>) => Promise<{ data: Float32Array; dims: number[] }>\n\nconst MODEL_ID = 'onnx-community/Qwen3-Embedding-0.6B-ONNX'\nconst CACHE_DIR = join(homedir(), '.qwenembedder', '.cache', 'models')\n\nlet pipe: PipeFn | null = null\nlet cache: Map<string, number[]> | null = null\nlet maxCacheSize = 100\n\nfunction setCache(key: string, value: number[]): void {\n if (!cache) return\n if (maxCacheSize === 0) return\n if (cache.size >= maxCacheSize) {\n const oldest = cache.keys().next()\n if (!oldest.done) {\n cache.delete(oldest.value)\n }\n }\n cache.set(key, value)\n}\n\ntype InitCommand = { _cmd: 'init'; cacheSize?: number }\ntype EmbedCommand = { _cmd: 'embed'; text?: string; texts?: string[] }\n\nrunAsWorker(async (command: InitCommand | EmbedCommand) => {\n if (command._cmd === 'init') {\n maxCacheSize = command.cacheSize ?? 100\n cache = new Map()\n return\n }\n\n if (!pipe) {\n pipe = (await pipeline('feature-extraction', MODEL_ID, {\n cache_dir: CACHE_DIR,\n dtype: 'q8',\n })) as unknown as PipeFn\n }\n\n if (command._cmd === 'embed') {\n const texts = command.texts ?? (command.text ? [command.text] : [])\n\n if (texts.length === 0) {\n throw new TypeError('Input must be a non-empty string')\n }\n\n const results: (number[] | undefined)[] = new Array(texts.length)\n const uncachedByText = new Map<string, number[]>()\n const uncachedTexts: string[] = []\n\n for (let i = 0; i < texts.length; i++) {\n const t = texts[i]\n if (typeof t !== 'string' || t.length === 0) {\n throw new TypeError('Input must be a non-empty string')\n }\n const cached = cache?.get(t)\n if (cached !== undefined) {\n results[i] = cached.slice()\n } else {\n const existing = uncachedByText.get(t)\n if (existing !== undefined) {\n existing.push(i)\n } else {\n uncachedByText.set(t, [i])\n uncachedTexts.push(t)\n }\n }\n }\n\n if (uncachedTexts.length > 0) {\n const output = await pipe!(uncachedTexts, { pooling: 'last_token', normalize: true })\n const data = Array.from(output.data) as number[]\n const hiddenDim = output.dims[output.dims.length - 1]\n\n for (let j = 0; j < uncachedTexts.length; j++) {\n const text = uncachedTexts[j]\n const vector = data.slice(j * hiddenDim, (j + 1) * hiddenDim)\n const stored = vector.slice()\n setCache(text, stored)\n const positions = uncachedByText.get(text) ?? []\n for (const position of positions) {\n results[position] = stored.slice()\n }\n }\n }\n\n if (command.text !== undefined) {\n return results[0] as number[]\n }\n return results as number[][]\n }\n})\n"],"mappings":";AAAA,SAAS,gBAAgB;AACzB,SAAS,mBAAmB;AAC5B,SAAS,eAAe;AACxB,SAAS,YAAY;AAIrB,IAAM,WAAW;AACjB,IAAM,YAAY,KAAK,QAAQ,GAAG,iBAAiB,UAAU,QAAQ;AAErE,IAAI,OAAsB;AAC1B,IAAI,QAAsC;AAC1C,IAAI,eAAe;AAEnB,SAAS,SAAS,KAAa,OAAuB;AACpD,MAAI,CAAC,MAAO;AACZ,MAAI,iBAAiB,EAAG;AACxB,MAAI,MAAM,QAAQ,cAAc;AAC9B,UAAM,SAAS,MAAM,KAAK,EAAE,KAAK;AACjC,QAAI,CAAC,OAAO,MAAM;AAChB,YAAM,OAAO,OAAO,KAAK;AAAA,IAC3B;AAAA,EACF;AACA,QAAM,IAAI,KAAK,KAAK;AACtB;AAKA,YAAY,OAAO,YAAwC;AACzD,MAAI,QAAQ,SAAS,QAAQ;AAC3B,mBAAe,QAAQ,aAAa;AACpC,YAAQ,oBAAI,IAAI;AAChB;AAAA,EACF;AAEA,MAAI,CAAC,MAAM;AACT,WAAQ,MAAM,SAAS,sBAAsB,UAAU;AAAA,MACrD,WAAW;AAAA,MACX,OAAO;AAAA,IACT,CAAC;AAAA,EACH;AAEA,MAAI,QAAQ,SAAS,SAAS;AAC5B,UAAM,QAAQ,QAAQ,UAAU,QAAQ,OAAO,CAAC,QAAQ,IAAI,IAAI,CAAC;AAEjE,QAAI,MAAM,WAAW,GAAG;AACtB,YAAM,IAAI,UAAU,kCAAkC;AAAA,IACxD;AAEA,UAAM,UAAoC,IAAI,MAAM,MAAM,MAAM;AAChE,UAAM,iBAAiB,oBAAI,IAAsB;AACjD,UAAM,gBAA0B,CAAC;AAEjC,aAAS,IAAI,GAAG,IAAI,MAAM,QAAQ,KAAK;AACrC,YAAM,IAAI,MAAM,CAAC;AACjB,UAAI,OAAO,MAAM,YAAY,EAAE,WAAW,GAAG;AAC3C,cAAM,IAAI,UAAU,kCAAkC;AAAA,MACxD;AACA,YAAM,SAAS,OAAO,IAAI,CAAC;AAC3B,UAAI,WAAW,QAAW;AACxB,gBAAQ,CAAC,IAAI,OAAO,MAAM;AAAA,MAC5B,OAAO;AACL,cAAM,WAAW,eAAe,IAAI,CAAC;AACrC,YAAI,aAAa,QAAW;AAC1B,mBAAS,KAAK,CAAC;AAAA,QACjB,OAAO;AACL,yBAAe,IAAI,GAAG,CAAC,CAAC,CAAC;AACzB,wBAAc,KAAK,CAAC;AAAA,QACtB;AAAA,MACF;AAAA,IACF;AAEA,QAAI,cAAc,SAAS,GAAG;AAC5B,YAAM,SAAS,MAAM,KAAM,eAAe,EAAE,SAAS,cAAc,WAAW,KAAK,CAAC;AACpF,YAAM,OAAO,MAAM,KAAK,OAAO,IAAI;AACnC,YAAM,YAAY,OAAO,KAAK,OAAO,KAAK,SAAS,CAAC;AAEpD,eAAS,IAAI,GAAG,IAAI,cAAc,QAAQ,KAAK;AAC7C,cAAM,OAAO,cAAc,CAAC;AAC5B,cAAM,SAAS,KAAK,MAAM,IAAI,YAAY,IAAI,KAAK,SAAS;AAC5D,cAAM,SAAS,OAAO,MAAM;AAC5B,iBAAS,MAAM,MAAM;AACrB,cAAM,YAAY,eAAe,IAAI,IAAI,KAAK,CAAC;AAC/C,mBAAW,YAAY,WAAW;AAChC,kBAAQ,QAAQ,IAAI,OAAO,MAAM;AAAA,QACnC;AAAA,MACF;AAAA,IACF;AAEA,QAAI,QAAQ,SAAS,QAAW;AAC9B,aAAO,QAAQ,CAAC;AAAA,IAClB;AACA,WAAO;AAAA,EACT;AACF,CAAC;","names":[]}
|
package/package.json
CHANGED
|
@@ -1,7 +1,7 @@
|
|
|
1
1
|
{
|
|
2
2
|
"name": "qwen-embedder",
|
|
3
|
-
"version": "0.1.
|
|
4
|
-
"description": "
|
|
3
|
+
"version": "0.1.4",
|
|
4
|
+
"description": "A lightweight, optimized local text embedding generation using qwen model",
|
|
5
5
|
"type": "module",
|
|
6
6
|
"main": "./dist/index.js",
|
|
7
7
|
"types": "./dist/index.d.ts",
|
|
@@ -29,12 +29,14 @@
|
|
|
29
29
|
"nlp",
|
|
30
30
|
"text-embedding",
|
|
31
31
|
"local-ai",
|
|
32
|
-
"feature-extraction"
|
|
32
|
+
"feature-extraction",
|
|
33
|
+
"qwen-embedder"
|
|
33
34
|
],
|
|
34
35
|
"license": "MIT",
|
|
35
36
|
"dependencies": {
|
|
36
37
|
"@huggingface/transformers": "^4.2.0",
|
|
37
|
-
"p-queue": "^9.3.0"
|
|
38
|
+
"p-queue": "^9.3.0",
|
|
39
|
+
"synckit": "^0.11.13"
|
|
38
40
|
},
|
|
39
41
|
"devDependencies": {
|
|
40
42
|
"@types/node": "^25.9.3",
|
package/src/embedder.ts
CHANGED
|
@@ -3,6 +3,8 @@ import PQueue from 'p-queue'
|
|
|
3
3
|
import { homedir } from 'os'
|
|
4
4
|
import { join } from 'path'
|
|
5
5
|
import type { EmbedderOptions } from './types.js'
|
|
6
|
+
import { SyncEmbedder } from './sync-embedder.js'
|
|
7
|
+
import type { SyncEmbedderOptions } from './sync-embedder.js'
|
|
6
8
|
|
|
7
9
|
const MODEL_ID = 'onnx-community/Qwen3-Embedding-0.6B-ONNX'
|
|
8
10
|
const CACHE_DIR = join(homedir(), '.qwenembedder', '.cache', 'models')
|
|
@@ -36,6 +38,10 @@ export class Embedder {
|
|
|
36
38
|
this.queue = concurrency !== null ? new PQueue({ concurrency }) : null
|
|
37
39
|
}
|
|
38
40
|
|
|
41
|
+
static createSync(options: SyncEmbedderOptions = {}): SyncEmbedder {
|
|
42
|
+
return new SyncEmbedder(options)
|
|
43
|
+
}
|
|
44
|
+
|
|
39
45
|
static async create(options: EmbedderOptions = {}): Promise<Embedder> {
|
|
40
46
|
const cacheSize = options.cacheSize ?? 100
|
|
41
47
|
const concurrency = resolveConcurrency(options)
|
|
@@ -78,7 +84,8 @@ export class Embedder {
|
|
|
78
84
|
}
|
|
79
85
|
|
|
80
86
|
const results: (number[] | undefined)[] = new Array(texts.length)
|
|
81
|
-
const
|
|
87
|
+
const uncachedByText = new Map<string, number[]>()
|
|
88
|
+
const uncachedTexts: string[] = []
|
|
82
89
|
|
|
83
90
|
for (let i = 0; i < texts.length; i++) {
|
|
84
91
|
const t = texts[i]
|
|
@@ -87,23 +94,33 @@ export class Embedder {
|
|
|
87
94
|
}
|
|
88
95
|
const cached = this.cache.get(t)
|
|
89
96
|
if (cached !== undefined) {
|
|
90
|
-
results[i] = cached
|
|
97
|
+
results[i] = cached.slice()
|
|
91
98
|
} else {
|
|
92
|
-
|
|
99
|
+
const existing = uncachedByText.get(t)
|
|
100
|
+
if (existing !== undefined) {
|
|
101
|
+
existing.push(i)
|
|
102
|
+
} else {
|
|
103
|
+
uncachedByText.set(t, [i])
|
|
104
|
+
uncachedTexts.push(t)
|
|
105
|
+
}
|
|
93
106
|
}
|
|
94
107
|
}
|
|
95
108
|
|
|
96
|
-
if (
|
|
109
|
+
if (uncachedTexts.length === 0) {
|
|
97
110
|
return results as number[][]
|
|
98
111
|
}
|
|
99
112
|
|
|
100
|
-
const uncachedTexts = uncached.map(i => texts[i])
|
|
101
113
|
const inferred = await this.runInference(uncachedTexts)
|
|
102
114
|
|
|
103
|
-
for (let j = 0; j <
|
|
115
|
+
for (let j = 0; j < uncachedTexts.length; j++) {
|
|
116
|
+
const text = uncachedTexts[j]
|
|
104
117
|
const vector = inferred[j]
|
|
105
|
-
|
|
106
|
-
|
|
118
|
+
const stored = vector.slice()
|
|
119
|
+
this.setCache(text, stored)
|
|
120
|
+
const positions = uncachedByText.get(text) ?? []
|
|
121
|
+
for (const position of positions) {
|
|
122
|
+
results[position] = stored.slice()
|
|
123
|
+
}
|
|
107
124
|
}
|
|
108
125
|
|
|
109
126
|
return results as number[][]
|
|
@@ -112,7 +129,7 @@ export class Embedder {
|
|
|
112
129
|
private async runInference(text: string): Promise<number[]>
|
|
113
130
|
private async runInference(texts: string[]): Promise<number[][]>
|
|
114
131
|
private async runInference(input: string | string[]): Promise<number[] | number[][]> {
|
|
115
|
-
const exec = () => this.pipe(input, { pooling: '
|
|
132
|
+
const exec = () => this.pipe(input, { pooling: 'last_token', normalize: true })
|
|
116
133
|
|
|
117
134
|
const output = this.queue
|
|
118
135
|
? await this.queue.add(exec)
|
|
@@ -135,12 +152,13 @@ export class Embedder {
|
|
|
135
152
|
private setCache(key: string, value: number[]): void {
|
|
136
153
|
if (this.maxCacheSize === 0) return
|
|
137
154
|
|
|
155
|
+
const stored = value.slice()
|
|
138
156
|
if (this.cache.size >= this.maxCacheSize) {
|
|
139
157
|
const oldest = this.cache.keys().next()
|
|
140
158
|
if (!oldest.done) {
|
|
141
159
|
this.cache.delete(oldest.value)
|
|
142
160
|
}
|
|
143
161
|
}
|
|
144
|
-
this.cache.set(key,
|
|
162
|
+
this.cache.set(key, stored)
|
|
145
163
|
}
|
|
146
164
|
}
|
package/src/index.ts
CHANGED
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
import { createSyncFn } from 'synckit'
|
|
2
|
+
import { fileURLToPath } from 'url'
|
|
3
|
+
import { dirname, resolve } from 'path'
|
|
4
|
+
|
|
5
|
+
const __dirname = dirname(fileURLToPath(import.meta.url))
|
|
6
|
+
const workerPath = resolve(__dirname, 'sync-worker.js')
|
|
7
|
+
|
|
8
|
+
type SyncFn = {
|
|
9
|
+
(command: { _cmd: 'init'; cacheSize?: number }): void
|
|
10
|
+
(command: { _cmd: 'embed'; text: string }): number[]
|
|
11
|
+
(command: { _cmd: 'embed'; texts: string[] }): number[][]
|
|
12
|
+
}
|
|
13
|
+
|
|
14
|
+
export interface SyncEmbedderOptions {
|
|
15
|
+
cacheSize?: number
|
|
16
|
+
}
|
|
17
|
+
|
|
18
|
+
export class SyncEmbedder {
|
|
19
|
+
private fn: SyncFn
|
|
20
|
+
|
|
21
|
+
constructor(options: SyncEmbedderOptions = {}) {
|
|
22
|
+
this.fn = createSyncFn(workerPath) as SyncFn
|
|
23
|
+
this.fn({ _cmd: 'init', cacheSize: options.cacheSize ?? 100 })
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
embedSync(text: string): number[]
|
|
27
|
+
embedSync(texts: string[]): number[][]
|
|
28
|
+
embedSync(input: string | string[]): number[] | number[][] {
|
|
29
|
+
if (Array.isArray(input)) {
|
|
30
|
+
return this.fn({ _cmd: 'embed', texts: input })
|
|
31
|
+
}
|
|
32
|
+
return this.fn({ _cmd: 'embed', text: input })
|
|
33
|
+
}
|
|
34
|
+
}
|
|
@@ -0,0 +1,96 @@
|
|
|
1
|
+
import { pipeline } from '@huggingface/transformers'
|
|
2
|
+
import { runAsWorker } from 'synckit'
|
|
3
|
+
import { homedir } from 'os'
|
|
4
|
+
import { join } from 'path'
|
|
5
|
+
|
|
6
|
+
type PipeFn = (texts: string[], options?: Record<string, unknown>) => Promise<{ data: Float32Array; dims: number[] }>
|
|
7
|
+
|
|
8
|
+
const MODEL_ID = 'onnx-community/Qwen3-Embedding-0.6B-ONNX'
|
|
9
|
+
const CACHE_DIR = join(homedir(), '.qwenembedder', '.cache', 'models')
|
|
10
|
+
|
|
11
|
+
let pipe: PipeFn | null = null
|
|
12
|
+
let cache: Map<string, number[]> | null = null
|
|
13
|
+
let maxCacheSize = 100
|
|
14
|
+
|
|
15
|
+
function setCache(key: string, value: number[]): void {
|
|
16
|
+
if (!cache) return
|
|
17
|
+
if (maxCacheSize === 0) return
|
|
18
|
+
if (cache.size >= maxCacheSize) {
|
|
19
|
+
const oldest = cache.keys().next()
|
|
20
|
+
if (!oldest.done) {
|
|
21
|
+
cache.delete(oldest.value)
|
|
22
|
+
}
|
|
23
|
+
}
|
|
24
|
+
cache.set(key, value)
|
|
25
|
+
}
|
|
26
|
+
|
|
27
|
+
type InitCommand = { _cmd: 'init'; cacheSize?: number }
|
|
28
|
+
type EmbedCommand = { _cmd: 'embed'; text?: string; texts?: string[] }
|
|
29
|
+
|
|
30
|
+
runAsWorker(async (command: InitCommand | EmbedCommand) => {
|
|
31
|
+
if (command._cmd === 'init') {
|
|
32
|
+
maxCacheSize = command.cacheSize ?? 100
|
|
33
|
+
cache = new Map()
|
|
34
|
+
return
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
if (!pipe) {
|
|
38
|
+
pipe = (await pipeline('feature-extraction', MODEL_ID, {
|
|
39
|
+
cache_dir: CACHE_DIR,
|
|
40
|
+
dtype: 'q8',
|
|
41
|
+
})) as unknown as PipeFn
|
|
42
|
+
}
|
|
43
|
+
|
|
44
|
+
if (command._cmd === 'embed') {
|
|
45
|
+
const texts = command.texts ?? (command.text ? [command.text] : [])
|
|
46
|
+
|
|
47
|
+
if (texts.length === 0) {
|
|
48
|
+
throw new TypeError('Input must be a non-empty string')
|
|
49
|
+
}
|
|
50
|
+
|
|
51
|
+
const results: (number[] | undefined)[] = new Array(texts.length)
|
|
52
|
+
const uncachedByText = new Map<string, number[]>()
|
|
53
|
+
const uncachedTexts: string[] = []
|
|
54
|
+
|
|
55
|
+
for (let i = 0; i < texts.length; i++) {
|
|
56
|
+
const t = texts[i]
|
|
57
|
+
if (typeof t !== 'string' || t.length === 0) {
|
|
58
|
+
throw new TypeError('Input must be a non-empty string')
|
|
59
|
+
}
|
|
60
|
+
const cached = cache?.get(t)
|
|
61
|
+
if (cached !== undefined) {
|
|
62
|
+
results[i] = cached.slice()
|
|
63
|
+
} else {
|
|
64
|
+
const existing = uncachedByText.get(t)
|
|
65
|
+
if (existing !== undefined) {
|
|
66
|
+
existing.push(i)
|
|
67
|
+
} else {
|
|
68
|
+
uncachedByText.set(t, [i])
|
|
69
|
+
uncachedTexts.push(t)
|
|
70
|
+
}
|
|
71
|
+
}
|
|
72
|
+
}
|
|
73
|
+
|
|
74
|
+
if (uncachedTexts.length > 0) {
|
|
75
|
+
const output = await pipe!(uncachedTexts, { pooling: 'last_token', normalize: true })
|
|
76
|
+
const data = Array.from(output.data) as number[]
|
|
77
|
+
const hiddenDim = output.dims[output.dims.length - 1]
|
|
78
|
+
|
|
79
|
+
for (let j = 0; j < uncachedTexts.length; j++) {
|
|
80
|
+
const text = uncachedTexts[j]
|
|
81
|
+
const vector = data.slice(j * hiddenDim, (j + 1) * hiddenDim)
|
|
82
|
+
const stored = vector.slice()
|
|
83
|
+
setCache(text, stored)
|
|
84
|
+
const positions = uncachedByText.get(text) ?? []
|
|
85
|
+
for (const position of positions) {
|
|
86
|
+
results[position] = stored.slice()
|
|
87
|
+
}
|
|
88
|
+
}
|
|
89
|
+
}
|
|
90
|
+
|
|
91
|
+
if (command.text !== undefined) {
|
|
92
|
+
return results[0] as number[]
|
|
93
|
+
}
|
|
94
|
+
return results as number[][]
|
|
95
|
+
}
|
|
96
|
+
})
|