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[][]> { async _embed(texts: string[]): Promise<number[][]> {
// replace newlines, which can negatively affect performance. // replace newlines, which can negatively affect performance.
const clean = texts.map((text) => text.replace(/\n/g, ' ')) const clean = texts.map((text) => text.replace(/\n/g, ' '))
const hf = new HfInference(this.apiKey)
const obj: any = { const obj: any = {
inputs: clean inputs: clean
} }
if (!this.endpoint) obj.model = this.model if (this.endpoint) {
return this.caller.call(() => this.client.featureExtraction(obj)) as Promise<number[][]> hf.endpoint(this.endpoint)
} else {
obj.model = this.model
}
const res = await this.caller.callWithOptions({}, hf.featureExtraction.bind(hf), obj)
return res as number[][]
} }
embedQuery(document: string): Promise<number[]> { async embedQuery(document: string): Promise<number[]> {
return this._embed([document]).then((embeddings) => embeddings[0]) const res = await this._embed([document])
return res[0]
} }
embedDocuments(documents: string[]): Promise<number[][]> { async embedDocuments(documents: string[]): Promise<number[][]> {
return this._embed(documents) return this._embed(documents)
} }
} }