huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-12 04:44:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix parsing

Browse Source
This commit is contained in:
Nicolas Patry 2024-09-06 18:50:52 +02:00
parent 785c6e4893
commit 69c168d1e0
No known key found for this signature in database
GPG Key ID: 64AF4752B2967863
1 changed files with 10 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
integration-tests/conftest.py
View File

@ -102,15 +102,16 @@ class ResponseComparator(JSONSnapshotExtension):
def _convert_data(data): def _convert_data(data):
if isinstance(data, Dict): if isinstance(data, Dict):
choices = data["choices"] if "choices" in data:
if isinstance(choices, List) and len(choices) >= 1: choices = data["choices"]
if "delta" in choices[0]: if isinstance(choices, List) and len(choices) >= 1:
return ChatCompletionChunk(**data) if "delta" in choices[0]:
if "text" in choices[0]: return ChatCompletionChunk(**data)
return Completion(**data) if "text" in choices[0]:
return ChatComplete(**data) return Completion(**data)
if isinstance(data, Dict): return ChatComplete(**data)
return Response(**data) else:
return Response(**data)
if isinstance(data, List): if isinstance(data, List):
return [_convert_data(d) for d in data] return [_convert_data(d) for d in data]
raise NotImplementedError raise NotImplementedError