Spaces:
Runtime error
Runtime error
File size: 1,172 Bytes
c93058d f1b879c c93058d f1b879c c93058d f1b879c c93058d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
import { pipeline } from "@xenova/transformers";
import { Embeddings, EmbeddingsParams } from "langchain/embeddings/base";
export interface XenovaTransformersEmbeddingsParams extends EmbeddingsParams {
model?: string;
}
export default class XenovaTransformersEmbeddings
extends Embeddings
implements XenovaTransformersEmbeddingsParams
{
model: string;
client: any;
constructor(fields?: XenovaTransformersEmbeddingsParams) {
super(fields ?? {});
this.model = fields?.model ?? "Xenova/all-MiniLM-L6-v2";
}
async _embed(texts: string[]): Promise<number[][]> {
if (!this.client) {
this.client = await pipeline("embeddings", this.model);
}
console.log(this.client, texts);
return this.caller.call(async () => {
return await Promise.all(
texts.map(async (t) => (await this.client(t, { pooling: 'mean', normalize: true })).data)
);
});
}
embedQuery(document: string): Promise<number[]> {
return this._embed([document]).then((embeddings) => embeddings[0]);
}
embedDocuments(documents: string[]): Promise<number[][]> {
return this._embed(documents);
}
} |