add fix to HF endpoint

This commit is contained in:
Henry 2023-07-08 21:31:43 +01:00
parent 6a7e378baf
commit 3cda348e95
3 changed files with 21 additions and 14 deletions

View File

@ -78,9 +78,7 @@ export class HuggingFaceInference extends LLM implements HFInput {
async _call(prompt: string, options: this['ParsedCallOptions']): Promise<string> {
const { HfInference } = await HuggingFaceInference.imports()
const hf = new HfInference(this.apiKey)
if (this.endpoint) hf.endpoint(this.endpoint)
const res = await this.caller.callWithOptions({ signal: options.signal }, hf.textGeneration.bind(hf), {
model: this.model,
const obj: any = {
parameters: {
// make it behave similar to openai, returning only the generated text
return_full_text: false,
@ -91,7 +89,13 @@ export class HuggingFaceInference extends LLM implements HFInput {
repetition_penalty: this.frequencyPenalty
},
inputs: prompt
})
}
if (this.endpoint) {
hf.endpoint(this.endpoint)
} else {
obj.model = this.model
}
const res = await this.caller.callWithOptions({ signal: options.signal }, hf.textGeneration.bind(hf), obj)
return res.generated_text
}

View File

@ -30,12 +30,11 @@ 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, ' '))
return this.caller.call(() =>
this.client.featureExtraction({
model: this.model,
inputs: clean
})
) as Promise<number[][]>
const obj: any = {
inputs: clean
}
if (!this.endpoint) obj.model = this.model
return this.caller.call(() => this.client.featureExtraction(obj)) as Promise<number[][]>
}
embedQuery(document: string): Promise<number[]> {

View File

@ -78,9 +78,7 @@ export class HuggingFaceInference extends LLM implements HFInput {
async _call(prompt: string, options: this['ParsedCallOptions']): Promise<string> {
const { HfInference } = await HuggingFaceInference.imports()
const hf = new HfInference(this.apiKey)
if (this.endpoint) hf.endpoint(this.endpoint)
const res = await this.caller.callWithOptions({ signal: options.signal }, hf.textGeneration.bind(hf), {
model: this.model,
const obj: any = {
parameters: {
// make it behave similar to openai, returning only the generated text
return_full_text: false,
@ -91,7 +89,13 @@ export class HuggingFaceInference extends LLM implements HFInput {
repetition_penalty: this.frequencyPenalty
},
inputs: prompt
})
}
if (this.endpoint) {
hf.endpoint(this.endpoint)
} else {
obj.model = this.model
}
const res = await this.caller.callWithOptions({ signal: options.signal }, hf.textGeneration.bind(hf), obj)
return res.generated_text
}