Ben Burtenshaw commited on
Commit
a212c60
1 Parent(s): 2373135

error handling in inference

Browse files
Files changed (1) hide show
  1. infer.py +8 -4
infer.py CHANGED
@@ -6,13 +6,17 @@ API_URL = (
6
  )
7
 
8
 
9
-
10
-
11
-
12
  def query(question, hub_token: str):
13
  payload = {
14
  "inputs": question,
 
 
 
 
15
  }
16
  headers = {"Authorization": f"Bearer {hub_token}"}
17
  response = requests.post(API_URL, headers=headers, json=payload)
18
- return response.json()[0]["generated_text"]
 
 
 
 
6
  )
7
 
8
 
 
 
 
9
  def query(question, hub_token: str):
10
  payload = {
11
  "inputs": question,
12
+ "parameters": {
13
+ "wait_for_model": True,
14
+ "return_full_text": False,
15
+ },
16
  }
17
  headers = {"Authorization": f"Bearer {hub_token}"}
18
  response = requests.post(API_URL, headers=headers, json=payload)
19
+ try:
20
+ return response.json()[0]["generated_text"]
21
+ except Exception:
22
+ return "Error occurred while querying the model."