add fix to HF Embeddings

This commit is contained in:
Henry 2023-07-20 01:07:29 +01:00
parent 491c4be8af
commit 741b97e80e
1 changed files with 13 additions and 5 deletions

View File

@ -30,18 +30,26 @@ export class HuggingFaceInferenceEmbeddings extends Embeddings implements Huggin
async _embed(texts: string[]): Promise<number[][]> {
// replace newlines, which can negatively affect performance.
const clean = texts.map((text) => text.replace(/\n/g, ' '))
const hf = new HfInference(this.apiKey)
const obj: any = {
inputs: clean
}
if (!this.endpoint) obj.model = this.model
return this.caller.call(() => this.client.featureExtraction(obj)) as Promise<number[][]>
if (this.endpoint) {
hf.endpoint(this.endpoint)
} else {
obj.model = this.model
}
embedQuery(document: string): Promise<number[]> {
return this._embed([document]).then((embeddings) => embeddings[0])
const res = await this.caller.callWithOptions({}, hf.featureExtraction.bind(hf), obj)
return res as number[][]
}
embedDocuments(documents: string[]): Promise<number[][]> {
async embedQuery(document: string): Promise<number[]> {
const res = await this._embed([document])
return res[0]
}
async embedDocuments(documents: string[]): Promise<number[][]> {
return this._embed(documents)
}
}