Merge branch 'main' into v0.24-release

huggingface · Jul 18, 2024 · 83b5d7f · 83b5d7f
2 parents 7fe286b + 1f9dbf8
commit 83b5d7f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
@@ -2251,7 +2251,12 @@ def text_generation(
  if stream:
  return _stream_text_generation_response(bytes_output, details) # type: ignore
 
- data = _bytes_to_dict(bytes_output)[0] # type: ignore[arg-type]
+ data = _bytes_to_dict(bytes_output) # type: ignore[arg-type]
+
+ # Data can be a single element (dict) or an iterable of dicts where we select the first element of.
+ if isinstance(data, list):
+ data = data[0]
+
  return TextGenerationOutput.parse_obj_as_instance(data) if details else data["generated_text"]
 
  def text_to_image(

diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py
@@ -2282,7 +2282,12 @@ async def text_generation(
  if stream:
  return _async_stream_text_generation_response(bytes_output, details) # type: ignore
 
- data = _bytes_to_dict(bytes_output)[0] # type: ignore[arg-type]
+ data = _bytes_to_dict(bytes_output) # type: ignore[arg-type]
+
+ # Data can be a single element (dict) or an iterable of dicts where we select the first element of.
+ if isinstance(data, list):
+ data = data[0]
+
  return TextGenerationOutput.parse_obj_as_instance(data) if details else data["generated_text"]
 
  async def text_to_image(