Harrison/hf inf error (#1543)

Co-authored-by: Konstantin Hebenstreit <57603012+KonstantinHebenstreit@users.noreply.github.com>
This commit is contained in:
Harrison Chase 2023-03-08 20:53:46 -08:00 committed by GitHub
parent 357d808484
commit 9405af6919
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -121,12 +121,15 @@ class HuggingFaceEndpoint(LLM, BaseModel):
)
except requests.exceptions.RequestException as e: # This is the correct syntax
raise ValueError(f"Error raised by inference endpoint: {e}")
generated_text = response.json()
if "error" in generated_text:
raise ValueError(
f"Error raised by inference API: {generated_text['error']}"
)
if self.task == "text-generation":
# Text generation return includes the starter text.
generated_text = response.json()
text = generated_text[0]["generated_text"][len(prompt) :]
elif self.task == "text2text-generation":
generated_text = response.json()
text = generated_text[0]["generated_text"]
else:
raise ValueError(