Spaces:
Sleeping
Sleeping
Joshua Lochner
commited on
Enable lazy construction of pipeline (#554)
Browse files
src/lib/server/websearch/sentenceSimilarity.ts
CHANGED
|
@@ -1,4 +1,4 @@
|
|
| 1 |
-
import type { Tensor } from "@xenova/transformers";
|
| 2 |
import { pipeline, dot } from "@xenova/transformers";
|
| 3 |
|
| 4 |
// see here: https://github.com/nmslib/hnswlib/blob/359b2ba87358224963986f709e593d799064ace6/README.md?plain=1#L34
|
|
@@ -6,8 +6,18 @@ function innerProduct(tensor1: Tensor, tensor2: Tensor) {
|
|
| 6 |
return 1.0 - dot(tensor1.data, tensor2.data);
|
| 7 |
}
|
| 8 |
|
| 9 |
-
|
| 10 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 11 |
// see https://huggingface.co/thenlper/gte-small/blob/d8e2604cadbeeda029847d19759d219e0ce2e6d8/README.md?code=true#L2625
|
| 12 |
export const MAX_SEQ_LEN = 512 as const;
|
| 13 |
|
|
@@ -17,6 +27,8 @@ export async function findSimilarSentences(
|
|
| 17 |
{ topK = 5 }: { topK: number }
|
| 18 |
) {
|
| 19 |
const input = [query, ...sentences];
|
|
|
|
|
|
|
| 20 |
const output: Tensor = await extractor(input, { pooling: "mean", normalize: true });
|
| 21 |
|
| 22 |
const queryTensor: Tensor = output[0];
|
|
|
|
| 1 |
+
import type { Tensor, Pipeline } from "@xenova/transformers";
|
| 2 |
import { pipeline, dot } from "@xenova/transformers";
|
| 3 |
|
| 4 |
// see here: https://github.com/nmslib/hnswlib/blob/359b2ba87358224963986f709e593d799064ace6/README.md?plain=1#L34
|
|
|
|
| 6 |
return 1.0 - dot(tensor1.data, tensor2.data);
|
| 7 |
}
|
| 8 |
|
| 9 |
+
// Use the Singleton pattern to enable lazy construction of the pipeline.
|
| 10 |
+
class PipelineSingleton {
|
| 11 |
+
static modelId = "Xenova/gte-small";
|
| 12 |
+
static instance: Promise<Pipeline> | null = null;
|
| 13 |
+
static async getInstance() {
|
| 14 |
+
if (this.instance === null) {
|
| 15 |
+
this.instance = pipeline("feature-extraction", this.modelId);
|
| 16 |
+
}
|
| 17 |
+
return this.instance;
|
| 18 |
+
}
|
| 19 |
+
}
|
| 20 |
+
|
| 21 |
// see https://huggingface.co/thenlper/gte-small/blob/d8e2604cadbeeda029847d19759d219e0ce2e6d8/README.md?code=true#L2625
|
| 22 |
export const MAX_SEQ_LEN = 512 as const;
|
| 23 |
|
|
|
|
| 27 |
{ topK = 5 }: { topK: number }
|
| 28 |
) {
|
| 29 |
const input = [query, ...sentences];
|
| 30 |
+
|
| 31 |
+
const extractor = await PipelineSingleton.getInstance();
|
| 32 |
const output: Tensor = await extractor(input, { pooling: "mean", normalize: true });
|
| 33 |
|
| 34 |
const queryTensor: Tensor = output[0];
|