diff --git a/benchmarks/backend_request_func.py b/benchmarks/backend_request_func.py index 0f13c79a..ea70a1f4 100644 --- a/benchmarks/backend_request_func.py +++ b/benchmarks/backend_request_func.py @@ -219,7 +219,15 @@ async def async_request_deepspeed_mii( if response.status == 200: parsed_resp = await response.json() output.latency = time.perf_counter() - st - output.generated_text = parsed_resp["text"][0] + if "choices" in parsed_resp: + output.generated_text = parsed_resp["choices"][0][ + "text"] + elif "text" in parsed_resp: + output.generated_text = parsed_resp["text"][0] + else: + output.error = ("Unexpected response format: " + "neither 'choices' nor 'text' found") + output.success = False output.success = True else: output.error = response.reason or ""