|
|
@ -122,7 +122,7 @@ class MosaicML(LLM):
|
|
|
|
|
|
|
|
|
|
|
|
prompt = self._transform_prompt(prompt)
|
|
|
|
prompt = self._transform_prompt(prompt)
|
|
|
|
|
|
|
|
|
|
|
|
payload = {"input_strings": [prompt]}
|
|
|
|
payload = {"inputs": [prompt]}
|
|
|
|
payload.update(_model_kwargs)
|
|
|
|
payload.update(_model_kwargs)
|
|
|
|
payload.update(kwargs)
|
|
|
|
payload.update(kwargs)
|
|
|
|
|
|
|
|
|
|
|
@ -160,15 +160,16 @@ class MosaicML(LLM):
|
|
|
|
# The inference API has changed a couple of times, so we add some handling
|
|
|
|
# The inference API has changed a couple of times, so we add some handling
|
|
|
|
# to be robust to multiple response formats.
|
|
|
|
# to be robust to multiple response formats.
|
|
|
|
if isinstance(parsed_response, dict):
|
|
|
|
if isinstance(parsed_response, dict):
|
|
|
|
if "data" in parsed_response:
|
|
|
|
output_keys = ["data", "output", "outputs"]
|
|
|
|
output_item = parsed_response["data"]
|
|
|
|
for key in output_keys:
|
|
|
|
elif "output" in parsed_response:
|
|
|
|
if key in parsed_response:
|
|
|
|
output_item = parsed_response["output"]
|
|
|
|
output_item = parsed_response[key]
|
|
|
|
|
|
|
|
break
|
|
|
|
else:
|
|
|
|
else:
|
|
|
|
raise ValueError(
|
|
|
|
raise ValueError(
|
|
|
|
f"No key data or output in response: {parsed_response}"
|
|
|
|
f"No valid key ({', '.join(output_keys)}) in response:"
|
|
|
|
|
|
|
|
f" {parsed_response}"
|
|
|
|
)
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
if isinstance(output_item, list):
|
|
|
|
if isinstance(output_item, list):
|
|
|
|
text = output_item[0]
|
|
|
|
text = output_item[0]
|
|
|
|
else:
|
|
|
|
else:
|
|
|
|