error handling for huggingface endpoints

This commit is contained in:
Krrish Dholakia 2023-08-14 10:21:57 -07:00
parent 017c4d5d03
commit fb07d6b389
6 changed files with 7 additions and 5 deletions

View file

@ -315,6 +315,8 @@ def completion(
response = requests.post(API_URL, headers=headers, json=input_payload)
## LOGGING
logging(model=model, input=prompt, custom_llm_provider=custom_llm_provider, additional_args={"max_tokens": max_tokens, "original_response": response.text}, logger_fn=logger_fn)
if isinstance(response, dict) and "error" in response:
raise Exception(response["error"])
completion_response = response.json()[0]['generated_text']
prompt_tokens = len(encoding.encode(prompt))
completion_tokens = len(encoding.encode(completion_response))